Learning World Models for Unconstrained Goal Navigation
Authors: Yuanlin Duan
Wensen Mao
He Zhu
Venue: NeurIPS 2024 poster (2024)
Keywords: World Models
Reinforcement Learning
Goal-Conditioned Reinforcement Learning
Model-Based Reinforcement Learning
Exploration Strategies
Primary area: reinforcement_learning

Abstract: Learning world models offers a promising avenue for goal-conditioned reinforcement learning with sparse rewards. By allowing agents to plan actions or exploratory goals without direct interaction with the environment, world models enhance exploration efficiency. The quality of a world model hinges on the richness of data stored in the agent's replay buffer, with expectations of reasonable generalization across the state space surrounding recorded trajectories. However, challenges arise in generalizing learned world models to state transitions backward along recorded trajectories or between states across different trajectories, hindering their ability to accurately model real-world dynamics. To address these challenges, we introduce a novel goal-directed exploration algorithm, MUN (short for "World Models for Unconstrained Goal Navigation"). This algorithm is capable of modeling state transitions between arbitrary subgoal states in the replay buffer, thereby facilitating the learning of policies to navigate between any "key" states. Experimental results demonstrate that MUN strengthens the reliability of world models and significantly improves the policy's capacity to generalize across new goal settings.

Learning World Models for Unconstrained Goal Navigation
Yuanlin Duan Rutgers University yuanlin duanOrutgers.cm edu
Wensen Mao
He Zhu Rutgers University hz3750cs ruttees..
Rutgers University eus00ecs.rutgerss
Abstract
Learning world models offers a promising avenue for goal-conditioned reinforce- ment learning with sparse rewards. By allowing agents plan actions ex- ploratory goals without direct interaction with the environment, world models enhance exploration efficiency. The quality of world model hinges the rich- ness of data stored in the agent's replay buffer, with expectations reasonable generalization across the state space surrounding recorded trajectories. However, challenges arise in generalizing learned world models to state transitions back- ward along recorded trajectories hindering their ability to accurately model real-world dynamics. To address these challenges, we introduce a novel goal-directed exploration algorithm, MUN (short for "World Models for Unconstrained Goal Navigation") "). This algorithm is capable of modeling state transitions between arbitrary subgoal states in the replay buffer, thereby facilitating the learning of policies to navigate between any "key" states. Experimental results demonstrate that MUN strengthens the reliability of world models and significantly improves the policy's 's capacity to generalize across new goal settings.
between states across different trajectories,
Introduction
Goal-conditioned reinforcement learning (GCRL) has emerged powerful framework for learning diverse skills within an environment and subsequently solving tasks based on user-specified goal commands, without requiring further training (Mendonca al. Andrychowie al., 2017). Given that specifying dense task rewards for GCRL requires domain expertise, access to object positions, is consuming, and is prone to human errors, rewards in GCRL are typically sparse, signaling success only upon reaching goal states. However, sparse rewards pose challenge for exploration during training. address this challenge, several previous methods, e.g., (Hafner et 2019a Hansen et al. 2023 Mendonca et al. 2021) have proposed learning generative world model the environment using a reconstruction (decoder) objective, an instantiation of Model-basd Reinforcement Learning (MBRL), visualized in Fig. [ This approach is appealing because the world model can provide a rich learning signal (Yu et al. 2020 Georgiev et al. 2024). ). For example, world models allow agents to plan their actions exploratory goals without directly interacting with the
real environment for more efficient exploration (Hu et al. 2023 Sekar et al. 2020).
Existing MBRL techniques train world models to capture the dynamics of the environment from the agent's past experiences stored a replay buffer. The richness of the data stored in the agent's 's replay buffer directly impacts the quality of World Model. It is expected that the world model generalizes reasonably well the state space surrounding the trajectories recorded in the replay buffer. However, the world model may not generalize well to state transitions backward along recorded trajectories or to states across different trajectories, which impedes the world model's ' learning of the real-world dynamics.
38th Conference on Neural Information Processing Systems (NcurIPS 2024).

---

To induce a data-rich replay buffer covering a wide range of dynamic transitions, in this paper, we present a novel goal-directed exploration al- gorithm for effective world model- ing and policy learning, MUN (short for "World Models for Unconstrained Goal Navigation"). MUN facilitates modeling state transitions between any subgoal states in the replay buffer, whether tracing back along recorded trajectories or transitioning between states separate trajectories. This enhances the reliability of the learned world model and signifi- cantly improves the generalizability of the policy derived from the model to real-world environments, thereby boosting the exploration capabilities of the method. Additionally, we introduce a simple and practical strategy for discovering key subgoal states from the replay buffer. The key subgoals precisely mark the milestones necessary for task completion, such as steps like grasping and releas- ing blocks in the context of block-stacking scenarios. By world modeling and policy learning for unconstrained navigation between these key states, MUN can generalize to new goal settings, such as block unstacking that was not given to the agent at training time.
Replay Ruffer
(1,)
World Model Mi
actor
(2:
o. (e 2: 145))
critic
Real World
Sample trajacteties
Figure 1: The general framework of model-based RI..
follows. First, propose novel goal-directed exploration algorithm MUN for effective world modeling of state transition between arbitrary subgoal states in replay buffers. As the quality of the world model improves, MUN becomes highly effective learning goal-conditioned policies that excel at exploration in sparse-reward environments. Second, we present practical strategy for identifying pivotal subgoal states, which serve as milestones in completing sophisticated tasks. By training world models for unconstrained transition between these milestones, our method enables learning policies that can adapt novel goal scenarios. Finally, we evaluate MUN in challenging robotics environments, such as guiding a multi-legged ant robot through a maze, maneuvering robot arm amidst cluttered tabletop objects, and rotating items in the grasp of an anthropomorphic robotic hand. Across these environments, MUN exhibits superior efficiency in training generalizable goal-cone condiionit policies compared baseline methods and ablations.
Our key contributions
2 Problem Setup and Background
We consider the problem of goal-conditioned reinforcement learning (GCRL) under a Markov Decision Process (MDP) parameterized by (5, A, P. G. n, R. Pu). S and are the state and action spaces, respectively. The probability distribution of the initial states is given by Pu(s), and P(s'|s, |s, a) the transition probability. 77 assumes that every state can be mapped to corresponding achieved goal g. The reward function R defined as R(s, a, 8', (' = 1 (7()²) = g). We assume that each episode has a fixed horizon For ease of presentation, we further assume S G and 1) is an identify function in this paper.
G a mapping from the state space to the goal space, which
S G A + IR+, which gives rise to trajectory samples of the form = { 40, g, 81,..., ST). The purpose of the policy is to learn how reach the goals drawn from the goal distribution Py- With a discount factor y (0, 1),
goal-conditioned policy is probability distribution
maximizes ((x) (ggp))
To- 7t R(st, at, $141,9)
In the context of model-based reinforcement learning (MBRL), a world model M trained over trajectories sampled from the agent's interactions with the real environment, which are stored in replay buffer, to predict the dynamics the real environment. Fig. Iillustrates the general MBRL framework. We use the world model structure recurrent state-space model (RSSM). We provide detailed
Dreamer (Hafnr et al. 2019a 2020 2023) to
learn real environment dynamics explanation of the network architecture and working principles of the RSSM in Appendix A.I Our study focuses tackling the world model learning problem in goal-conditioned model based reinforcement learning settings. Particularly, we consider GC-Dreamer (goal-conditioned Dreamer) as important baseline with the following learning components:
World Model: (1-1)))
Actor: T (a.ls..))
Critic: V(s..)
(I)

---

Dak rapie
Aeber Depperpen te reci
Woeld Model
Actes:
Meck
yov Vinkk
Ardier Onipprr Cares bokk redesbi
World Moxid
(a) Key subgoal states a 3-Block Stacking task.
(b) Bidirectional Replay Buffer
Figure 2: In Fig- 2(a) we illustrate the key states involved in completing the task of 3-block stacking. In Fig. 2(b) we demonstrate the significant advantages of the bidirectional replay buffer used in MUN over traditional methods in learning world models.
In GC-Dreamer, the goal coal-conditioned agent (als, g) samples goal commands g G from the given environment goal distribution Pg to collect trajectories in the real world. These trajectories are used to train the world model M and subsequently, n is trained on imagined rollouts generated by M using the model-based actor-critic algorithm in Dreamer (Hafner et al. 2020), with these two steps run in alternation. The critic estimates the sum of future rewards >, r?, and the actor tries to maximize the predicted values from the critic. The goal goal-reaching reward y. is defined by the self-supervised temporal distance network D, (Mendonca al. 2021), i.e. (s,g)) -D: Ds(s, D; predicts the anticipated number of action steps needed to transition from s to g. Essentially, x is reinforced to minimize the action steps required transition from the current state to a sampled goal state g. The temporal distance estimator D, is trained by extracting pairs of states and s from an imagined rollout generated by running the policy over the world model and predicting the distance & between them as follows:
Dj (V (s(+k)) k/H
Here, T represents the preprocessing for imagined states, such as transforming them into the world model's latent space (we assume G in the paper). H represents the total length of the imagined rollout. Further details on the training procedure of D, can be found in Appendix A.2
3 Training World Models for Unconstrained Goal Navigation
In this section, we introduce MUN, our main approach to addressing the core challenge in GCRL: ef- ficient exploration long-horizon, sparse-reward environments. Our approach focuses on enhancing the agent's understanding of the real-world environment through improved dynamic (world) modeling and latent space representation. As the quality of the world model improves, the goal-conditioned policy developed from it generalizes more effectively to the real environment. By closing the gen- cralization gap between the policy's s behavior in the real environment and the world model, MUN effectively guides the agent's 's exploration towards the desired goal region in the real environment.
3.1 Training Generalizable World Models
Fig illustrates the general framework of Model-hased RI., where world models are trained us- ing agent's experiences stored in a replay buffer populated with observed environment transitions (St, az, St+1) linking the environment's future states 91+1 and past states along with the corre- sponding control actions a, The richness of the environment space and dynamic transitions captured by the replay buffer define the extent of what a world model can learn about the real environment. Through supervised learning, the model can generalize reasonably well within the state space moving forward along the trajectories recorded in the replay buffer. However, it may be inaccurate for the state transitions moving backward along the recorded trajectories or across different trajectories. Consider the task of stacking blocks using a robot manipulator in Fig. 2(a) When humans learn to stack blocks, they also understand how to reverse the process to unstack the blocks or return to the initial state. In contrast, world model trained solely on data from policies under training for stacking
3

---

Algorithm The main training framework of MUN
1: Input: Policy T., World Model M, reward function y subgoals transfer number Na, subgoal time limit T,
2: Initialize buffers D, DDAD, Dege
3: for = 1 to Ntrain do
4: if Should Plan Subgoals then
5: Bege A batch of episodes from Dege
6:
Geubgoats f- DDAD(Bege) with Algorithm]
7:
Initialize empty trajectory 7
8:
for to N, do
9:
4,
10:
9a Sample a subgoal randomly from G Guugoaats while agent has not reached ys and igg T, do
11:
12:
Append one step in real environment with σ using goal to т
to to
DAAD DDAD (T}
15:
I' f Trajectory of πh sampled using the environment goal distribution Pg
Dege t Dege U
17:
D D+DDAD U ege
18:
Update M with D
Update T in imagination with M maximize G
unlikely to accurately model the unstacking process. As a result, the model may yield hallucinated trajectories for training policies, causing significant discrepancy between the policy's "slicyss behavior in the model and in the real world, thereby leading to ineffective exploration.
To improve model generalizability, in MUN, we proposed to learn world models capable of char- acterizing state transitions between any states in the replay buffer, whether by tracing back along recorded trajectories or transitioning between states on separate trajectories. Fig, 2b visualizes the comparison between the bidirectional replay buffer for learning world models used MUN and the unidirectional replay buffer in conventional model-bard based algorithms. The bidirectional replay buffer not only covers a wider observation space but also captures richer set of dynamic transitions. As discussed in Sec. [ due to joint optimization, the richer set of dynamic transitions in MUN allows for more reliable latent representation of the environmental space and consequently higher quality reward function (Equation 2 for training policies generalizable to the real environment on top of the learned model.
We depict the learning algorithm in MUN in Algorithm] | In the algorithm, we maintain Gsuhgoals set of pivot subgoal states sampled from the relay butfer (illustrated in Algorithm 2 and aim to learn world models capable of seamless transitions between these subgoals. At line we periodically update G Guubools as the training evolves. In the loop starting from line 8, we repeatedly sample N. subgoals from Subgoals and direct the agent to sequentially reach these subgoals within a time limit of T, steps for each. In this way, MUN samples replay buffer that records bidirectional state transitions between the subgoals in G subguals- Based on our experience, we find that setting N, = 2 is sufficient. Further discussion on the setting of N₂ provided in Sec 4 At line we train the world model M using trajectories collected by both the goal commands from Gsudgoats (stored DDAD) and that sampled from the environment goal distribution Pg (stored in Dege). Then, we sample imaginary rollouts from the world model for policy training at line
Comparison with Go-Explore. We highlight the key difference between MUN's ss exploration strategy and the recently popular "Go-Explore" strategy (Ecoffet al., 2019 Pislar et al. 2021] Tuyls et al.) 2022, Hu et al. 2023), designed for spploration-extensie each training episode comprises two phases: the "Go-phase" and the "Explore explorr--hasee" During the "Go-phase," the goal-conditioned policy directs the agent to an "interesting" goal (e.g., states with low frequency of occurrence in the replay buffer) (Pong al. 2019 Pitis et al. 2020), resulting in final state ST, after T, steps. Following this, the "Explore-phase begins, where an undirected exploration policy takes over from 8T, for the remaining T, timesteps. This exploration policy trained to maximize an intrinsic exploration reward (c.g., to visit areas of the real world that the World
long- term GCRL settings. In Go-Explore,

---

Model has not yet learned well) (Bellemare et al., 2016 Pathak et al. 2017 Burda et al. 2018 Sekar al., 2020). This structure of training episodes has been shown result in richer exploration (Pislar et al. 2021). In MUN, when N, 2, A Algorithm | essentially replaces the "Explore-phase in "Go-Explore" with another "Go-phase". Thus, the algorithm directs the agent to navigate between two "interesting" goals selected from the replay buffer. Firstly, MUN is computationally efficient as it eliminates the need to train a separate exploration policy and an ensemble of world models used to generate intrinsic exploration rewards. Secondly, MUN trains the world model for unconstrained navigation between goal states in the replay buffer, thereby improving the model's generalization to the real-world environment and leveraging the model for exploration. We empirically compare the two strategies in the context of model-based GCRL in Sec.
3.2 Key Subgoal Generation through Distinct Action Discovery (DAD)
Having set up the main learning algorithm, we seek to address: how should we pick exploration- inducing goals from the replay buffer at training time to help learn generalizable world models? A straightforward strategy is to sample trajectories from the replay buffer and select subgoal states at fixed time intervals along these trajectories. We improve this simple approach with a practical method called DAD (Distinet Action Discovery) for identifying key subgoal states, which represent the pivotal milestones necessary to complete a complex task. Consider the block stacking task as an example. The robotic arm must be able move its gripper to the vicinity of block, close the gripper. lift the block, move to the top of another block, release the block, and finally open the gripper. These key subgoal states are essential for completing the task. We illustrate the roles of key states in 3-Block Stacking task in Fig. 2(a) For an agent to learn this task, it must master reaching and navigating between the key states. By training world models for unconstrained transitions between these key states, MUN can develop models that more accurately capture the task structure and learn policies capable of adapting to novel goal scenarios e.g. block unstacking.
There exist methods for identifying key states (Zhang al.] 2021] Paul et al. 2019). However, these methods often tend to be overly complex, leading to insuffi- cient generalization across different envi- ronments and requiring adjustments the methods' components or parameters for various tasks. Our approach is based on the observation that certain actions are crucial at different stages of task completion. For instance, in a block stacking task, the robotic arm must learn actions such as closing the gripper, grasping the block, lifting it, and releasing the gripper. When the agent performs these key actions, the corresponding states can often considered key subgoal states. By selecting actions that significantly differ along trajectories and extracting the corresponding states during these actions, we can identify potential key subgoal states. The Farthest Point Sampling (FPS) algorithm [Eldar et al. 1997) provides a simple and efficient method for selecting N points with maximal differences from set. We apply FPS to choose N time steps with the greatest variations in actions from a batch of trajectory data, thereby obtaining the set of key subgoal states corresponding to these time steps. Algorithm 2 shows how MUN finds key subgoal states using the DAD method.
Algorithm 2 Key Subgoal Generation by Distinet Action Discovery
1: Function: DAD(....
2: Input: A batch of episodes Bege, number of subguals N Nasagoula
3: A Find the most different Nrubgoots actions from Bage by
FPS
Sautyoots get the corresponding states of A from Bage
5: Guubyools ((SSuungua)s)
6: rettur Gauhgolla
Experiments
We evaluate MUN across various robotic manipulation and navigation environments, aiming to address the following three research questions: (RQ1) Does MUN outperform other goal-conditioned model-based reinforcement learning baselines with advanced exploration strategies? (RQ2) Can DAD effectively identify key subgoal states along trajectories to the environment goal region? (RQ3) Does MUN successfully leverage the bi-directional replay buffer to train a generalizable policy for navigating effectively between arbitrary subgoals?
4.1 Environments
We conducted experiments on six challenging goal-conditioned tasks to evaluate MUN. Ant-Maze, an aik-like robot is tasked to learn complex 4-legged locomotion behavior and navigate around the
5

---

Ant Maze
Walker
3-Block Stacking
Pen Rotation
Block Rotation
Fetch Slide
Figure 3: We evaluate MUN on 6 environments: Ant Maze, Walker, 3-Block Stacking, Block Rotation,
Pen Rotation, Fetch Slide.
hallways within a maze structure. The Walker task involves two-legged robot learning to control its leg joints effectively to achieve stable walking to reach goals along a flat plane forward or backward. In 3-Block Stacking, robot arm with two-fingered gripper operates on tabletop with three blocks. The goal is to stack the blocks into a tower configuration. The agent needs to learn to push, pick, and stack objects while discovering complex action sequences to complete the task in the environment. Previous solutions have relied on methods like demonstrations, curriculum learning, or extensive simulator data, highlighting the task's 's difficulty (Ecoffet et al., 2019: Li et al. [2000) Nair et al.] 2018 Lanier 2019). The Block Rotation and Pen Rotation tasks require the agent to manipulate a block and a pen, respectively, to achieve randomly specified orientation along all axes. Pen Rotation is particularly challenging due to the pen's thinness, requiring precise control to prevent it from dropping. In Fetch Slide, manipulator slides puck designated goal on a slippery table. Unlike tasks that involve direct manipulation, Fetch Slide emphasizes the challenge of accurately controlling the force and direction of the push operation, as the puck must slide across the flat surface to the target. See Appendix. C for more information about environments.
4.2 Baselines
We compare MUN with the following baselines. The GC-Dreamer baseline is discussed in Sec. We include two baselines based on the Go-Explore strategy (Ecoffet et al. 2019) that has been proved efficient in the GCRL setting: MEGA (Pitis et al., 2020) and PEG (Hu et al. (223)) A Go-Explore agent firstly uses its goal-conditioned policy πC to approach a sampled exploration-inducing goal command g. referred as the Go-phase. In the Explore-phar paplore-phaee,..m. it activates an exploration policy я to explore the environment from the terminal state of the Go-phase. In contrast, MUN improves the generalization of world models facilitate effective real-world environment exploration. During training, MUN collects trajectories that navigate between two goal states sampled from its candidate in (co-Exploree with another "Go-phase"
subgoal set, essentially replacing the MEGA commands the agent to rarely seen states at the frontier by using kernel density estimates (KDE) of state densities and chooses low-density goals from the replay buffer. PEG selects goal commands to guide an agent's s goal-conditioned policy toward states with the highest exploration potential given its current level of training. This potential is defined as the expected accumulated exploration reward during the Explore-phase Similar to MUN, our baseline methods, named PEG-G and MEGA-G, augment GC-Dreamer with the PEG and MEGA Go-Explore strategies, respectively. In these methods, the replay buffer D contains not only trajectories sampled by the GCRL policy ma commanded by environment goals but also exploratory trajectories sampled using the Go-Explore strategies. The exploration policy N E in PEG-G and MEGA-G G is the Plan2Explore policy from Sekar cet al. (2020), which encourages the agent to actively search for states that induce disparities among an ensemble of world models.
We note that MUN and the baselines are all implemented based on the Dreamer framework as realized in GC-Dreamer
4.3 Results
Fig. shows the evaluation performance of MUN and all baselines across training. MUN demonstrates superior performance compared to the baseline models, excelling in both the final success rate and the speed of learning. MUN outperforms the Go- Go-Explore baselines (MEGA-G and PEG-G) across all tasks, demonstrating the effeectivnnes of the exploration strategy in MUN over the alternative Go- Go-Explore strategies. In the most challenging tasks block stacking, block rotation, and pen rotation-MuL shows a significant margin of superiority. For example, MUN achieves over 95%
MUN is not tied specific world model architecture and can be applied to any model-based RL framework.

---

1.90
MUN
160
HLR PBG REGA-O
HEDA-
NBSA-G
CC wammer
GC DreaTeer
Creamer
3.40
8.00
340
8:0
(a) Ant Maze
(b) Walker
(c) 3-Block Stacking
3.0
340
VEGA-C
VEGA-G
HEGG-G OC Dieamer
0C Greamer
GC Dicamo
xcepa
sdegs
(d) Block Rotation
(c) Pen Rotation
(f) Fetch Slide
Figure 4: Experiment results comparing MUN with the baselines over 5 random seeds.
success rate block stacking, while all other baselines only manage to achieve around 60% success rate on this task within 2.5M steps. MEGA-G and PEG-G heuristically pick exploration-inde -nntuiiinn goals initiate exploration by separate policy. Since finding goal state that is optimally aligned with both the goal-conditioned policy and the exploration policy is challenging, these methods can result in suboptimal goals, thereby slowing down exploration. GC-Dreamer lacks Go-Explore phase, which limits its exploration potential. Despite this, it can still perform comparably to or even better than MEGA-G and PEG-G in certain contexts. This indicates that the Go-Explore strategy does not always guarantee improved exploration, and suboptimal goal-setting during the "Go-phase" can hinder exploration (see 3 block stacking).
Fetch Slide is a non-prehensile manipu- lation task. This environment has asym- metric state transitions: when the puck is slid outside the robot's workspace, the ma- nipulator cannot reach the puck's position to slide it backward due to physical con- straints. MUN still outperforms the other baselines in this environment. We found MUN, with the DAD strategy, can discover key subgoals for this task, like contacting the puck, placing the manipulator at dif- ferent angles around the puck, and stabi- lizing the manipulator upon reaching the goal (these key states result from distinct actions). MUN enables learning state tran- sitions between these key subgoals to dis- cover the high-level task structure. It learns between any positions within the workspace and predicts low probabilities for infeasible transitions from puck positions outside the workspace. Particularly. it enables the agent to hit the puck multiple times if it is within its workspace, thereby improving task success rates. That said, the current goal selection mechanism in MUN lacks process to filter out infeasible goals from the current state, which could adversely affect sample efficiency. We left addressing this limitation and implementing robust filtering mechanism for infeasible goals focus for future work.
MUN PED-G MERAGG
HUN PCG- Drcanrer
(a) 3-Block Stacking
(b) Rotation
The world model prediction curves throughout the training steps for 3-Block Stacking and Pen Rotation.
Figure 5:
generalizable world model that handles sliding the puck
We studied the prediction error of learned world models MUN and the baselines. Fig. 5 shows the one-step model prediction error throughout the training steps. The world models trained by MUN
7

---

show much smaller generalization gap to the real environment compared to the baselines across the training steps. Consequently, MUN can effectively leverage these higher-quality world models to train policies that generalize better to the real environment. We present quantitative comparison of the world model prediction quality between MUN and the baselines in terms of model prediction compounding error in Appendix F3
4.4
Can DAD find key subgoals?
3-Block Stacking
Ant Maze
Walker
in three environments: Ant-Maze, Walker, 3-Block Stacking. They present the important landmarks on the path to the task goal regions.
Figure 6: Key subgoals found by DAD (Algorithm 2
We visualize several subgoals found by the DAD algorithm during the training process in Fig. for three environments: Ant-Maze, Walker, 3-Block Stacking. In Walker, DAD successfully identifies crucial joint angles and forces of the Walker robot during its forward locomotion, including standing, striding, jumping, landing, and leg support. In Ant-Maze, DAD recognizes significant motion variations at corridor corner. In 3-Block Stacking. DAD successfully identifies crucial state transitions required during the stacking process. These critical subgoals include block grasping, lifting, horizontal movement, vertical movement, and gripper release. For more discussion about subgoals found by the DAD in other environments, please refer to Appendix E.
4.5
Can MUN navigate between arbitrary subgoals?
80.00%
60.00%
MUN
GC-Dreamer
MEGA-G G
48 10.00%
PEG-G
20.00%
0.00%
Figure 7: Experiment setup and results of navigation between any pair of subgoals in the 3-Block Stacking environment. In the left part, the bottom section of each image depiets the ultimate evaluation goal for one evaluation episode, while the top section illustrates the manually set initial state. The right part shows the evaluation success rates.
As MUN is capable of identifying pivotal subgoal states necessary for complex tasks and training world models and policies for seamless transitions between these subgoals, investigate MUN's capacity to generalize to new task settings concerning important subgoals. We set the initial state of the agent at one random subgoal and command it to reach another random subgoal. Such task setting is not provided to the agent during training. For the 3 Block Stacking task, we employ set of 15 manually created subgoals representing various critical states in the block-stacking process, resulting in 225 unique combinations of initial states and test goals for evaluation. Each combination undergoes 10 repeated evaluations, totaling 2250 evaluation trajectories. These evaluations encompass both the forward and reverse processes of stacking and unstacking blocks, assessing the agent's proficiency in both task completion and restoration.
example, the left portion of Fig. 7 we visualize some subgoals used as initial task state in the upper part and some subgoals used as evaluation test goals in the lower part. The right section of Fig. illlstraates MUN's superiority over the other baselines in these evaluation experiments, achieving the highest success rate through its ability to develop

---

1.6.
MUN
MUN TODAD
MUN RODAD
0.70
MUN-noDAA
MUN- SeyCbs
MUN- Mey-eeyyy
MUN- MeNN--eyybbs
MUNN-Na33
0.5
MUUN Na-33
6.44
MUN--is-3
0.00
6.0
0.90
0.:00
0.:00
0.0
0110
(a) 3-Block Stacking
(b) Block Rotation
(c) Pen Rotation
Figure Experiment results comparing MUN with its ablations over random seeds.
robust and adaptable world model that generalizes to novel tasks. Additional results in different
environments are provided in Appendix F2
4.6 Ablation study
We conducted the following ablation studies to investigate MUN's exploration goal selection mecha- nism. First, we investigated the effect of the number of subgoal states (N,) in our algorithm. MUN sequentially traverses N, 2 goal states sampled from the replay buffer explore the enviromment during each training episode. We introduced an ablation MUN-Ns-3 that sets N, = 3. This ablation aims to investigate whether increasing N₂ leads improved learning performance. Second, considered an ablated version of MUN, named MUN-noDAD, which replaces the goal sampling strategy DAD (Algorithm{) ( with simple method that chooses goal states with fixed time interval in trajectories sampled from the replay buffer. This ablation investigates the importance of identifying key subgoal states, which represent pivotal milestones necessary complete complex task. seeks to determine whether training world models from state transitions between these key states in MUN is essential, if using any states from the replay buffer would suffice. Lastly, we explored an alternative key subgoal discovery strategy. MUN identifies key subgoals for exploration as states in the replay buffer that result in distinct actions within the action space. We introduced an ablation, MUN-KeyObs which directly discovers key subgoals from the state space by identifying centroids of (latent) state clusters the replay buffer, following the strategy Zhang et al.] (2021).
The results are depicted in Fig. g MUN outperforms all ablated versions. Setting N = 3 slows down the training performance, supporting our claim suffices set N, 2. The performance of MUN-noDAD and MUN-KeyObs does not match MUN, especially in the 3 Block Stacking environment, highlighting that discovering key subgoals in the action space (the DAD strategy) indeed contributes to higher performance and efficiency. It is noteworthy that the ablation methods achieve relatively small gap success rates compared MUN in the challenging Block Rotation and Pen Rotation environments. This suggests that MUN's approach to learning a world model from state transitions between any states in the replay buffer (whether tracing back along recorded trajectories transitioning across separate trajectories) alone is effective in bridging the generalization gap between the model and the real environment.
Related Work
Model-based reinforcement learning (MBRL) is a promising approach to reinforcement learning that learns a model of the environment and uses it to plan actions (Sutton 1991] Deisenroth and
Rasmussen 201T Oh et al. 2017; Chua al. 2018 It has achieved remarkable success in numerous
control tasks and games, such as chess (Silver et al. 2017 Schrittwieser et al., 2020; Xu et al. 2022),
Atari games (Hafner et al.) 2020; Schrittwieser et [2000] Oh et al.. 201 continuous control tasks (Kurutach et al. [018 Buckman et al. [2018] Halner et al. 2019b] Janner et al., 2019), and robotic manipulation tasks (Lowrey et al.) 2018 Luo et al. 2018). 8). The dynamic model serves as pivotal component of model-based reinforcement learning, primarily fulfilling two key roles: planning actions (Deisenroth and Rasmussen] 2011 Oh et al.) 2017| Chua et al., 2018 Lowrey et al., 2018
Hafner al.) (019b) generating synthetic data to aid in the training of model-free reinforcement

---

learning algorithms (Jannrr et al. 2019 Hafner et al., 2020 2023). The primary drawback of the former lies in the excessive cost associated with long-term planning. To address this issue, the concept of ensemble has been employed to enhance performance (Chua et al. 2018 Kurutach et al. 2018 BBukkman al., 2018). Oh et al. (0177) and (ansen et al. (2022b) integrate the dynamics model with value prediction network to improve the accuracy of long-term planning. The latter also suffers from the potential bias of the model, which can result inaccuracies in the generated data, thereby
directly impacting policy learning (Luo et al. 2018 Lai et al. 2021
Multi- goal reinforcement learning (RL) agents (Schaul et al. 2015: Plappert al. 2018 Ghosh
et al., 2019) acquire goal-conditioned behaviors capable of achieving and generalizing across diverse sets objectives. Researchers have been continuously exploring the integration of Model-based RL and Goal Goal-conditioned RL (Mendonca et al. [02I]] Nair et al., 2020, /Zhang et al. (020), leveraging the capabilities of dynamic models in planning and generating synthetic data enhance the training efficiency and generalization of GCRL. However, compared to traditional RL problems, GCRL faces more severe challenges regarding reward sparsity and exploration difficulties (Ren et al.) 2019
Florensa et al. 20188 Trott et al. 2019). These challenges often lead to significant biases in the learned World Model, consequently impairing the performance of goal-conditioned policies (Mendonca et al. 2021 Hu et al. 2023). Pong et al. (219) propose to learn maximum-entropp goal distribution, Pitis et al. (020) encourage the agent to explore goals with low frequency of occurrence in the replay buffer. Sekar et al. (2020) introduce planning algorithm to pick goals for exploration using World Models.
World Models hold significant promise for GCRL,, as they enable fast exploration and support the training of more generalized policies (McCarthy et al., [202]] Shyam et al.) 2019 Hu et al. 2023]
Sekar et al.] 2020). Within this framework, learning a reliable World Model is essential for developing
effective policies Zhang et al. 2024: Young et al. 2022 Wang et al. 2023 Lai et al. 2021). 1). Kauvar et al. (0233) propose curiosity-driven exploration method, which is focused
replay buffer management. Hansen et al. (2022a) use demonstration data as a supplement to the replay buffer to learn more reliable World Model. Previous work has often focused on devising more appropriate objectives when sampling real trajectory data from the environment to enrich the diversity dynamic transitions in the replay buffer (Nair et al. 20200 Charlesworth and Montana, 22000] Trott et al.) 2011
Florensa et al., 2018; Campero et al. 2020). However, they overlooked the overall direction of dynamic transitions within the data which extremely affects the richness of dynamic transitions leam a comprehensive World Model.
Conclusion
In summary, we introduce MUN, novel goal-directed exploration algorithm designed for effective world modeling of seamless transitions between arbitrary states in replay buffers, whether retracing along recorded trajectories or transitioning between states on separate trajectories. As the quality of the world model improves, MUN demonstrates high efficacy in lcarning goal-conditioned policies in sparse-reward environments. Additionally, we present practical strategy DAD for identifying pivotal subgoal states, which act as critical milestones in completing complex tasks. The experimental results underscored the effectiveness of MUN in strengthening the reliability of world models and learning policies capable of adapting to novel test goals.
Reproducibility Statement
The code for MUN is available on https:// https/hhwb.mmmmmeittttttteeeeddd -aaetttddd-Reeeniinn ieesooniggGGroup MUN For hyperparameter settings and baseline pscudocode, please refer to Appendix D and Ap- pendix
Acknowledgements
We thank the anonymous reviewers for their comments and suggestions. This work was supported by NSF Award #CCF-2124155.
10

---

References
Andrychowicz, M., Wolski, F., Ray, A., Schneider, J., Fong, R., Welinder, P., McGrew, B., Tobin, J., Pieter Abbeel, O., and Zaremba, W. (2017). Hindsight experience replay. Advances in neural information processing systems, 30.
Bellemare, M., Srinivasan, S., Ostrovski, G., Schaul, T., Saxton, D., and Munos, R. (2016). Unifying count-based exploration and intrinsic motivation. Advances in neural information processing systems, 29.
Buckman, J., Hafner, D., Tucker, G., Brevdo, E., and Lee, H. (2018). Sample-efficient reinforcement learning with stochastic ensemble value expansion. Advances in neural information processing systems, 31.
Burda, Y., Edwards, H., Storkey, A., and Klimov, O. (2018). Exploration by random network distillation. arXiv preprint 188.94481..1284.
Campero, A., Raileanu, R., Küttler, H., Tenenbaum, J. B., Rocktäschel, T., and Grefenstette, E. (2020). Learning with amigo: Adversarially motivated intrinsic goals. arXiv preprint 1rXiv:2006.12122.
Charlesworth, H. and Montana, G. (2020). Plangan: Model-based planning with sparse rewards and multiple goals. Advances Neural Information Processing Systems, 33:8532-8542.
Chua, K., Calandra, R., McAllister, R., and Levine, S. (2018). Deep reinforcement learning in handful of trials using probabilistic dynamics models. Advances in neural information processing systems, 31.
Deisenroth, M. and Rasmussen, C. E. (2011). Pilco: A model-based and data-efficient approach to policy search. In Proceedings of the 28th International Conference on machine learning (ICML-II), pages 465-472.
Ecoffet, A., Huizinga, Lehman, J., Stanley, K. O., and Clune, (2019). Go-explore: new approach for hard-exploration problems. arXiv preprint arXiv:1901.1em
Eldar, Y. Lindenbaum, M., Porat, M., and Zeevi, Y. Y. (1997). The farthest point strategy for progressive image sampling. IEEE transactions on image processing, 1305--335
Florensa, C., Held, D., Geng, X., and Abbeel, P. (2018). Automatic goal generation for reinforcement learning agents. International conference on machine learning. pages 1515-1528. PMLR.
Georgiev, I., Giridhar, V., Hansen, N., and Garg, A. (2024). Pwm: Policy learning with large world models. arXiv preprint 2rXhv:2407.02460
Ghosh, D., Gupta, A., Reddy, A., Fu, J., Devin, C., Eysenbach, B., and Levine, S. (2019). Learning to reach goals via iterated supervised learning, arXiv preprint arXiv: 19/2.06088.
Hafner, D., Lillicrap, T., Ba, J., and Norouzi, M. (2019a). Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv:/912.001 1912.01603.
Hafner, D., Lillicrap, T., Fischer, I., Villegas, R., Ha, D., Lee, H., and Davidson, J. (2019b). Learning latent dynamics for planning from pixels. In International conference on machine learning, pages 2555 2565. PMLR.
Hafner, D., Lillicrap, T., Norouzi, M., and Ba, J. (2020). Mastering atari with discrete world models. arXiv preprint arXiv:2010.02/90
Hafner, D., Pasukonis, J., Ba, J., and Lillicrap, T. (2023). Mastering diverse domains through world models. arXiv preprint arXivv2301.04104.
Hansen, N., Lin, Y. Su, H., Wang, X., Kumar, V., and Rajeswaran, A. (2022a). Modem: Accelerating visual model-based reinforcement learning with demonstrations, arXiv preprint ar2.v2221..05698
Hansen, N., Su, H., and Wang, (2023). Td-mpe2: Scalable, robust world models for continuous control. arXiv preprint ar ar11:22110.1882.

---

Hansen, N., Wang, X., and Su, H. (2022b). Temporal difference learning for model predictive control. arXiv preprint arXhv:2203.04955
Hu, E. S., Chang, R., Rybkin, O., and Jayaraman, D. (2023). Planning goals for exploration. arXiv preprint ar 23Xhv:2303.13002.
Janner, M., Fu, J., Zhang. M., and Levine, S. (2019). When to trust your model: Model- baaeed policy optimization. Advances in neural information processing systems, 32.
Kauvar, I., Doyle, C., Zhou, L.,. and Haber, N. (2023). Curious replay for model-based adaptation. arXiv preprint arXiv:2306. 15934.
Kurutach, T., Clavera, I., Duan, Y., Tamar, A., and Abbecl, P. (2018). Model-ensemble trust-region policy optimization. arXiv preprint arXiv: 1802. 10592.
Lai, H., Shen, J., Zhang, W., Huang, Y., Zhang, X., Tang, R., Yu, Y., and Li, Z. (2021). On effective scheduling of model-based reinforcement learning. Advances in Neural Information Processing Systems, 34:3694-7705.
Lai, Y., Wang, W., Yang, Y., Zhu, J., and Kuang, M. (2020). Hindsight planner. In Proceedings of the 19th International Conference on Autonomous Agents and MultiAgent Systems, pages 690-698.
Lanier, J. B. (2019). Curiosity-driven multi-criteria hindsight experience replay. University of California, Irvine.
Li, R., Jabri, A., Darrell, T., and Agrawal, P. (2020). Towards practical multi-object manipulation using relational reinforcement learning. In 2020 ieee international conference on robotics and automation (icra), pages 4051-4058. IEEE.
Lowrey, K., Rajeswaran, A., Kakade, S., Todorov, E., and Mordatch, (2018). Plan online, learn of- fline: Efficient learning and exploration via model-based control. arXiv preprint arXiv: 101:.0111100000
Luo, Y., Xu, H., Li, Y., Tian, Y., Darrell, T., and Ma, (2018). Algorithmic framework for model- based deep reinforcement learning with theoretical guarantees. arXiv preprint arXiv: 1807.03858.
McCarthy, R., Wang. Q., and Redmond, S. J. (2021). Imaginary hindsight experience replay: Curious basel learning for sparse reward tasks. arXiv preprint wrXiv:2110002414
Mendonca, R., Rybkin, O., Daniilidis, K., Hafner, D., and Pathak, D. (2021). Discovering and achieving goals via world models. Advances in Neural Information Processing Systems, 34:24379- 24391.
Nagabandi, A., Konolige, K., Levine, S., and Kumar, V. (2020). Deep dynamics models for learning dexterous manipulation. In Conference on Robot Learning, pages 1101-1112. PMLR.
Nair, A., McGrew, B., Andrychowicz, M., Zaremba, W., and Abbeel, P. (2018). Overcoming exploration in reinforcement learning with demonstrations. In 2018 IEEE international conference on robotics and automation (ICRA), pages 6292-6222 6299.. IEEE.
Nair, S., Savarese, S., and Finn, C. (2020). Goal-aware prediction: Learning to model what matters.
In International Conference on Machine Learning, pages 7207-7219. PMLR.
Oh, J., Singh, S., and Lee, H. (2017). Value prediction network. Advances in neural information processing systems, 30.
Pathak, D., Agrawal, P., Efros, A. A., and Darrell, T. (2017). Curiosity-driven exploration by self-supervised prediction. In International conference on machine learning, pages 2778-2787. PMLR.
Paul, S., Vanbaar, J., and Roy-Chowdhury, A. (2019). Learning from trajectories via subgoal
discovery. Advances in Neural Information Processing Systems, 32.
Pislar, M., Szepesvari, D., Ostrovski, G., Borsa, D., and Schaul, T. (2021). When should agents
explore? arXiv preprint arXiv:2108.1123
12

---

Pitis, S., Chan, H.. Zhao, S., Stadie, B., and Ba, J. (2020). Maximum entropy gain exploration for long horizon multi- goal reinforcement learning. In International Conference on Machine Learning. pages 7750-7761. PMLR.
Plappert, M., Andrychowizz. M., Ray, A., McGrew, B., Baker, B., Powell, G., Schneider, J., Tobin, J., Chociej, M., Welinder, P., et al. (2018). Multi-goal reinforcement learning: Challenging robotics environments and request for research, arXiv preprint arXiv: 1802.09464.
Pong, V. H., Dalal, M., Lin, S., Nair, A., Bahl, S., and Levine, S. (2019). Skew-fit: State-covering self-supervised reinforcement learning. arXiv preprint arXiv: 00698.3988.
Ren, Z., Dong, K., Zhou, Y., Liu, Q. and Peng, J. (2019). Exploration via hindsight goal generation. Advances in Neural Information Processing Systems, 32.
Schaul, T., Horgan, D., Gregor, K., and Silver, D. (2015). Universal value function approximators. In International conference on machine learning, pages 1312 1320. PMLR.
Schrittwieser, J., Antonoglou, I., Hubert, T., Simonyan, K., Sifre, L., Schmitt, S., Guez, A., Lockhart, E.,, Hassabis, D., Graepel, T., et al. (2020). Mastering atari, go, chess and shogi by planning with a learned model. Nature, 588(7839).004-099
Sekar, R., Rybkin, O., Daniilidis, K., Abbeel, P., Hafner, D., and Pathak, (2020). Planning to explore via self-supervised world models. In International conference on machine learning, pages 8583-8592. PMLR.
Shyam, P., Jaskowski, W., and Gomez, F. (2019). Model-based active exploration. In International conference on machine learning, pages 5779-5788. PMLR.
Silver, D., Hubert, T., Schrittwieser, J., Antonoglou, I., Lai, M., Guez, A., Lanctot, M., Sifre, L., Kumaran, D., Graepel, T., et al. (2017). Mastering chess and shogi by self-play with a general reinforcement learning algorithm. arXiv preprint arXiv:
Sutton, R. S. (1991). Dyna, an integrated architecture for learning, planning, and reacting. ACM Sigart Bulletin, 2(4): 16061660.
Trott, A., Zheng, S., Xiong, C., and Socher, R. (2019). Keeping your distance: Solving sparse reward tasks using self-balancing shaped rewards. Advances in Neural Information Processing Systems, 32.
Tuyls, J., Yao, S., Kakade, S., and Narasimhan, K. (2022). Multi-stage episodic control for strategic exploration in text games. arXiv preprint arXiv:2201.01251.
Wang, X., Wongkamjan, W., Jia, R., and Huang, F. (2023). Live in the moment: Learning dynamics model adapted to evolving policy. In International Conference on Machine Learning, pages 36470-36493, PMLR.
Williams, G., Aldrich, A., and Theodorou, E. (2015). Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.01 149.
Xu, Y., Hansen, N., Wang, Z., Chan, Y.-C., Su, H., and Tu, Z. (2022). On the feasibility of cross-task transfer with model-based reinforcement learning. arXiv preprint arXiv:2210. 10763.
Young, K., Ramesh, A., Kirsch, L., and Schmidhuber, J. (2022). The benefits of model-based generalization in reinforcement learning. arXiv preprint arXiv:2211.02222.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. (2020). Mopo: Model-based offline policy optimization. Advances in Neural Information Processing Systems, 33:14129-14142.
Zhang, I.., Yang, G., and Stadie, B. C. (2021). World model as graph: Learning latent landmarks for planning. In International conference on machine learning, pages 12611-. PMLR.
Zhang, W., Wang, G., Sun, J.. Yuan, Y., and Huang, G. (2024). Storm: Efficient stochastic transformer based world models for reinforcement learning. Advances in Neural Information Processing Systems, 36.
Zhang, Y., Abbeel, P., and Pinto, I. (2020). Automatic curriculum learning through value disagree- ment. Advances in Neural Information Processing Systems, 33:7648-7659,
13

---

Appendix
A Extended Background
A.1 Dreamer World Model
The RSSM consists of an encoder, recurrent model, a representation model, a transition predictor, and a decoder, as formulated in Equation [ And it employs an end-to-end training methodology, where its parameters are jointly optimized based on the loss functions of various components, including dynamic transition prediction, reward prediction, and observation encoding-decoding. These components often operate in a latent space rather than the original observation space, as encoded by the World Model. Therefore, during end-to-end training, the losses of all components indirectly optimize the latent space.
The encoder fE encodes the input state into embed state which is then fed with the deter- ministic state h2 into the representation model qø to generate the posterior state 2,. The transition predictor P& predicts the prior state & based on the deterministic state h, without access to the current input state zt. Using the concatenation of either (h,, 2t) or (he, &t) input, the recurrent transition function f iteratively updates the deterministic state h2 with given action
Encoder:
c fa(e)(x))
Recurrent model:
hr = fo(he-1,
21-1, 1- ac-1)
Representation model:
2e as (zelter ei)
Transition predictor:
2e ~ pa(Eelle))
Decoder: 2. fD(eolhe, 21)
A.2 Temporal Distance Training in LEXA
The goal- reaching reward G is defined by the self-supervised temporal distance objective (Mendoncal al., 2021) which aims to minimize the number of action steps needed to transition from the current state to a goal state within imagined rollouts. We use by to denote the concatenate of the deterministic state h4 and the posterior state at time step t.
b (h2, 21)
The temporal distance D; is trained by sampling pairs of imagined states by, be+* from imagined rollouts and predicting the action steps number k between the embedding them, with a predicted embedding é, from b, to approximate the true embedding of the observation 2t.
Predicted embedding:
emb(bt)) = è, so where e, = fE(x))
(5)
Temporal distance: (, 2e k/H
where = emb(br) êcts emb(be4k)
(6)
(bb. = - Di(e,,e44))
B Limitations and Future Work
The MUN has provided powerful guidance in enhancing world model learning by repeatedly studying the transitions between various key states. This allows the acquisition of richer dynamic transitions and deepens the world model's understanding an efficient strategy for discovering key states, as evidenced by the comparative results of the MUN and MUN-noDAD. We found that although DAD excels discovering key states with its simple and efficient method, it will identify ineffective and task-irrelevant states in tasks with highly complex
the real world. However, such a framework requires
14

---

action spaces or weak correlations between goal space and action space. This can lead to the degradation of the MUN architecture due to poor-quality subgoals, resulting in a substantial amount of ineffective sampling in the environment. Therefore, for environments with more complex action and goal spaces, we need to develop a more robust and effective method for discovering subgoals than DAD. Only with an efficient and powerful self-superrvsed subgoal discovery mechanism can the MUN framework be fully utilized.
Meanwhile, MUN autonomously discovers subgoals and learns a more robust and comprehensive world model by randomly navigating between subgoals. Although the MUN has achieved huge success in model-based reinforcement learning (MBRL), we believe it can also be applied to general model-free methods. General model-free methods do not require learning a world model and have a simpler architecture. The MUN can directly guide the condiiooit policy to enhance learning in navigation between different subgoals. can use sampled trajectories to learn this policy directly, bypassing the use of the world model to train policies and value functions through simulated trajectories, thereby enhancing the agent's ability to reach unconstrained goals. Therefore, we plan to explore the application and effectiivness of the MUN in model-free RL in the future and develop a new robust self-super subgoal discovery mechanism to make the MUN applicable to more complex environments.
C
Environments
C.1 3-Block Stacking
In this task, the robot must stack three blocks in different colors into a tower shape. While PEG assesses goals of varying difficulty levels: 3 easy goals (picking up a single block), 6 medium goals (stacking two blocks), and hard goals (stacking three blocks), our evaluation is focused solely on the 6 hard goals, and we use only 3 hard goals of them as the guiding goals from the training environment. Training and evaluating with only the hardest goals imposes a significant challenge for the MUN. However, we observed that the MUN can spontaneously discover additional easy and medium goals through DAD, as these serve as critical transitional states toward the hard goals. The environment is characterized by a 14-dimensional state and goal space. The first five dimensions capture the gripper's state, while the remaining nine dimensions correspond to the zyz positions of each block. The action space 4-dimensional, with three dimensions dedicated to the gripper's zyz movements and the fourth dimension controlling the gripper's finger movement. Success is defined by achieving an L2 distance of less than 3 cm between each block's zyz position and its target position. This environment is modified version of the FetchStack3 environment from Pitis al. (2020), designed to better test the robot's precision in stacking.
C.2 Walker
In this environment, 2D walker robot is trained and evaluated on its ability to move across flat surface. The environment's 's implementation is based on the code from Mendonca et al. (202I). To thoroughly assess the agent's 's capability and precision in covering longer distances, we expanded the evaluation goals to 12 1113, +16, 119, 225, +28) along the x axis from the initial position. In training setup for the MUN, we only use the goals at 113 and 16 provided by the environment, but we evaluate the agent's 's performance across all 12 goals. Success is measured by verifying whether the agent's ' position is within small margin of the target x position. The state and goal space in this environment are dine-dimensional, comprising the walker's ' 22 positions and its joint angles. This configuration ensures a comprehensive evaluation of the walker's locomotion capabilities.
C.3 Ant Maze
This environment builds upon the Ant Maze from Pitis et al. (2020), incorporating a few modifi- cations. The state and goal spaces in the Ant Maze environment are highly complex, totaling 29 dimensions. These dimensions include the ant's ' zyz position, joint angles, and velocities. The first three dimensions account for the zyyz position, the next 12 dimensions capture the joint angles of the ant's limbs, and the remaining 14 dimensions represent the velocities of the joints and the ant's movements in the ry plane. The action space consists of 8 dimensions, controlling the hip and ankle movements of the ant's 's legs. We matched the goal space to the state space, which includes the ant's
15

---

ryz coordinates, joint positions, and velocities. We also introduced an additional room in the top left increase the difficulty like PEG. In this scenario, the ant robot must traverse from the bottom left to the top left of a maze, navigating through various corridors. The task is particularly challenging due to its lengthy duration-each episode lasts 500 timesteps- and the significant distance the ant must cover. Unlike PEG, which evaluates goals in both the central hallway and the top left room, our evaluation focuses exclusively on the four most difficult goals located in the top left room. For training, we utilize all 32 goals throughout the maze. The maze itself measures about by 8 meters. The ant succeeds if its zy position is within 1.0 meter of the goal, roughly the size of single cell in the maze.
C.4 Fetch Slide
In this task, a robotic arm with a two-fingered gripper must push an object along a flat surface to specific goal location. We use the "FetchSlide-v environment from Gymnasium, where the robot operates in 25-dimensional state space that includes the robot's joint states, object position, and goal information. The goal space is 3- 3-dimensional, representing the target coordinates for the object. Each episode presents a unique random goal location within a bounded area, requiring the agent adjust its pushing strategy accordingly. A key challenge in Fetch Slide lies in the indirect manipulation of the object. The agent must accurately control the force and direction of its push while accounting for physical properties like friction, surface irregularities, and object momentum. Unlike grasping or lifting tasks, sliding demands precise force calibration and anticipation of the object's response to contact. For evaluation, the agent's 's learned policy is tested across 50 episodes with different goal locations, assessing its ability to generalize over varied configurations. Training goals are randomly generated from the environment, helping the agent explore diverse sliding trajectories to improve robustness across different scenarios.
C.5 Block and Pen Rotation
In this task, a robotic hand must manipulate either a thin pen or a block to achieve specified target rotations. We use "HandManniplaaPPeRRRttaayW
and v1" versions of the gymnasium environments. Both tasks feature a state space of 61 dimensions, encompassing the robot's joint states, object states, and goal information. The goal space is dimensional, representing the target pose details. Each episode will have randomized target rotations goal for all axes of the block and for the x and y axes of the pen. The pen is more challenging to handle due to its tendency to slip, requiring more precise control compared to the block. For evaluation, the latest policy is tested 50 episodes for each task, with each episode having a unique random goal. In our framework, training goals are also randomly generated from the environment.
D Baselines
We first present our overall training framework for goal-conditioned model- reinforcement learning (MBRL). It is important note that all baselines utilize this training framework, differing only in the strategy employed for collecting trajectories within the real environment. Our training framework is based on the implementation of LEXA (Mendonca al. 202T.
Algorithm General MBRL Training Framework
1: Input: Policy G F, Environment Goal Distribution G, World Model 17, reward function J.C.
E
2: D < {} Initialize buffer.
3: for Episode i = 1 to Nirain do
4:
7 <<- Collect trajectories(...)
5:
D+DUT
6:
Update world model M with D
7:
Update G in imagination with AM to maximize G Update n in imagination with AM to maximize r
8:
16

---

D.1 Go-Explore
Our baselines utilize state-of-the-art Go-Explore exploration framework, following the imple- mentation detailed in PEG (Hu et al., 2023). This approach initially employs goal-conditioned policy T to get as close as possible to a specified goal g. process referred to as the "Go phase." Subsequently, explorer policy x² is used to further explore the environment starting from the final state the Go phase, known as the "Explore phase."
The quality of the trajectories generated by the Go-Explore strategy largely depends on the selection of the goal g during the Go phase. Therefore, establishing an effective mechanism for selecting the Go phase goals is crucial. If the chosen goal g is too simple, the explorer will not sufficiently explore environment. Conversely, the goal is too difficult, the goal-achieving policy π. will fail to approach it effectively. Thus, the baselines MEGA-G and PEG-G employ different goal selection strategies to determine g. guiding the agent to areas with high exploration potential during the Go phase. MEGA-G and PEG-G enhance the agent's exploration efficiency by crafting robust exploration strategics, enabling faster learning of the world model with respect to new dynamic transitions and environmental areas. We present the pseudocode for Go- Go-Explore in Algorithm 4
Algorithm Go Explore Framework
function GO-EXPLORE(g, x x,,T x)
2: 8g - env.reset()
3: T {so}
4:
for Step f to T do
5:
8t f env.step(m (st-1.9))
6:
- T (s {st)
7:
if agent reach g then break
8:
9:
Le =1
10:
for Step t = te to te + Explore do
11:
env.step(a (st-1))
12:
T+ TU
13:
return
D.2
GC-Dreamer
GC-Dreamer the goal-conditioned version of Dreamer (Hafner et al., 2019a b 2020), without incorporating any exploration goal-directed strategies. It only uses goal-conditioned policy to collect trajectories, with goals provided by the training environment.
Algorithm 5 GC-Dreamed Goal Sampling
function COLLECT TRAJECTORIES(..)
2: + Returned by environment
3:
Sample trajectories by G using goal g
4: return 7
D.3 PEG-G
PEG uses world model to simulate exploration trajectories and evaluates the exploration Potential(PE(( (g)) to identify areas worth exploring.
[F(g) -Epca
(c F [sT)]
(8)
T+Te
By
1
T=
17

---

PEG set goal for the coal-conditioned policy and generalize it to K trajectories using world model. ST denotes the final state of the goal-conditioned trajectory from the "Go phase" of Go- Go-Explore to reach the g. Since the objective in Equation g is not easily computable, state distribution induced by the target-conditioned policy :T' which may rapidly change throughout the training process, it's crucial to use the latest estimates for better exploration. PEG achieve this by leveraging the learned world model. PEG utilize the learned exploration value function VE(sE) (Equation I0 from the learned world model estimate the exploration value of the final state for each trajectory, and average these estimates.
it relies on the final
1 K
E (llr)) [V
Σ
(s
where s4 ((C((1a)()))
(10)
K
( = p(s))]
II
M 81 ( (0-15-1-
(11)
The goals sampled for evaluating this exploration potential metric in PEG are drawn from distribution updated by the MPPI method (Williams et al., [2015] Nagabandi et al. 2020). PEG-G not only uses goals obtained by optimizing Equation t0 guide exploration sampling but also directly samples trajectories using goal-conditioned policy with goals provided by the environment. The sampling alternate between these two strategies as shown in the pseudocode in Algorithm
Algorithm PEG-G Sampling
1: function COLLECT TRAJECTORIES(..) if episodde i%2 then
2:
3:
- Optimize Equation 10 with MPPI
4:
GO-EXPLORE
5: else
6:
9 Returned by environment
7:
+ Sample trajectories by G using goal g
8:
return
D.4 MEGA-G
MEGA Piiii et al. 2020) employs kernel density estimates (KDE) to assess state densities and selects goals with low densities from the replay buffer. For the implementation of MEGA, we adopt the model-based MEGA methodology described in the PEG paper without modifications. The PEG paper has illustrated that their adaptation of MEGA outperforms the original MEGA implementation. This entails integrating MEGA's 's KDE model and incorporating a goal-conditioned value function into the LEXA framework filter goals based on reachability. Similar to PEG-G, MEGA-G switches between utilizing goals from the environment and employing the MEGA goal selection strategy.
Algorithm 7 MEGA-G Goal Sampling
1: function COLLECT TRAJECTORIES(..) if episodd i%2 = then
2:
3:
mingeD F(g) <- GO-EXPLOREE B)
4:
5: else
6:
Returned by environment
7:
f Sample trajectories by π using goal g
8:
return 7
D.5 MUN-noDAD
consider the Time-sample Hindsight Waypoints Sampling Strategy from Hindsight Planner
(Lai
et al. 2020) alternative the subgoal selection mechanism in MUN. MUN-noDAD selects
18

---

subgoals at fixed time intervals along trajectories, providing a simple and effective strategy for defining subgoals. MUN-noDAD can still benefit from the framework of MUN in navigating between different subgoals. The pseudocode for this baseline is as follows:
Algorithm 8 MUN-noDAD Subgoal Picking Strategy
Function: Subgoals, fixed (ntedd__aevv)(...))
2: Input: A batch of episodes B Beger number of subgoals N eubgoale
3: Saudgoats f pick N Naubgouls states at fixed time intervals from Bege
4: Gsubgoals + (SSuubooal))
5: return G Gsubgools
Algorithm 9 Trainning Frame for MUN-noDAD
1: Input: Policy G.G, World Model M, reward function " subgoals transfer number N., subgoal
time limit T,
2: Initialize buffers D,DDAD, Dege
3: for = 1 to Ntrain do
4: if Should Plan Subgoals then
5: Bege + A batch of episodes from Dege
6:
Gsubguals Subgoals fixed interval(...) with Algorithm g
7:
Initialize empty trajectory T
8:
for ===10 to N₂ do
9:
t, =0
10:
9. Sample subgoal randomly from Guugooais
11:
while agent has not reached 9s and ts < Tg do Append one step in real environment with α using goal g,
13:
1, + La +1
14:
DDAD +- DDAD (T)
Dege f- DegelJ Sample a trajectory with @ using goal from training environment
16:
DE -DDAD Dege
17:
Update M with D
18:
Update n in imagination with M maximize G
E
Implementation Details
E.1
Farthest Point Sampling (FPS) Algorithm
Algorithm 10 Farthest Point Sampling (FPS)
1: function FPS(points, num samplss)
2: sampled_points +
3: first_point + random.choiceeooo sandom..hoice(points)
4: campled_points.eee wpp.pednnts.apeddddd(pp apppappeffss_ppiinnn (firttppoint)
5: min distanccesse + [float('inf')] x len(points)
6: for each point points do
7:
min lin__distanee(p] +- distance(p, first_point)
8:
for iteration i 1 to num Laatmm_smmleesss do farthest_ Loiittppontttttte index f- arimax(min distances) farthest_pooie + foints/farthest Locittpooitiitteeeeeeeeeeeee index] sampled (aantttapppcnn((attttsttpppp
9:
10:
11:
12:
for each point p in points do
13:
min Lin__dstances(p] <- min(min hatiiiiiinditteeespp.. distance(p. porthest_point))
14:
return sampled pointss
19

---

The pseudocode presented in the Algorithm 10 illlusttee the process of Farthest Point Selection (FPS) algorithm. The FPS algorithm begins by initializing an empty list called 'sampled_ppe fooiiddpoits''' to store the selected points. The process commences by randomly selecting an initial point from the input point set, denoted as 'points', and adding it to "sampled Looeedd_pontts' Subsequently, "min distances' is initialized to keep track of the minimum distance from each point to any of the sampled points, with initial values set to infinity.
The core procedure involves iteratively selecting points until the desired number of samples is reached. At each iteration, the algorithm identifies the point in 'points' with the maximum minimum distance to the previously sampled points and includes it in 'sampled points'. Concurrently, 'min distances' updated reflect the recalculated minimum distance of each point to any of the sampled points.
E,2 Runtime
Table Runtimes per experiment.
Total Runtime (Hours)
Total Steps
3-Block Stacking
70
2.5e6
Walker
40
1.5e6
Ant Maze
36
le6
Block Rotation
2.5e6
Pen Rotation
68
2.5e6
Fetch Slide
2e6
We conduct each experiment on GPU Nvidia A100 and require about 3GB of GPU memory. See table in Table specific running time of MUN for different task. Most of the runtime is consumed by the neural network updates for the policy and the world model, while the time taken by DAD to filter subgoals minimal.
E.3 Hyperparameters
We use the default hyperparameters of the LEXA backbone MBRL agent (e.g.. learning rate, optimizer, network architecture) and keep them consistent across all baselines. MUN primarily requires hyperparameter tuning in the following: 1) the number of candidate subgoals stored N Nuuhgoatsi 2) the number of subgoals used for navigation when sampling in the environment N₂; and 3) the total episode length and the maximum number of timesteps allocated for navigating to a specific subgoal T.. We show these hyperparameters in Table
Table 2: Hyperparameters of MUN.
N Nsubgoats
N,
L
T
3-Block Stacking
20
2
150
75
Walker
10
2
150
75
Ant Maze
20
2
500
250
Block Rotation
20
2
150
75
Pen Rotation
20
2
150
75
Fetch Slide
20
2
150
75
F Additional Experiments
F.1 More subgoals found by DAD
We visualize several subgoals found by the DAD algorithm during the training process in Fig. 9 In Walker, the first five images show that DAD successfully identifies the crucial joint angles and forces of the Walker robot during its forward locomotion, including standing, striding, jumping, landing,
20

---

Figure 9: More subgoals found by DAD(Algorithm in all six environments
and leg support. In the subsequent three images, DAD similarly succeeds in recognizing the key movements of the Walker robot during its backward locomotion. In Ant-Maze, DAD recognizes significant motion variations corridor corners. In Block Rotation and Pen Rotation, DAD is able to identify crucial finger movements subgoals for rotating objects. In 3-Block Stacking, DAD successfully identifies crucial state transitions required during the stacking process. These critical subgoals include block grasping, lifting, horizontal movement, vertical movement, and gripper release.
F.2 Navigation Experiments
We the extend navigation experiments on 3-Block Stacking, Ant Maze, and Walker environments to if the MUN can learn better world model to navigate to unconstrained goals from unconstrained start state compared to other basclines. In the 3-Block Stacking task, we use a set 15 goals that represent various critical states in the block-stacking process. These goals serve candidates for both initial states and endpoint goals, resulting in a total of 225 unique combinations of initial states and endpoint goals for each evaluation episode. For each combination, evaluations,
conduct 10 repeated goal is ultimately computing the average success rate across 2250 evaluation trajectories. Our whether MUN effectively achieve random goal when the agent starts from an arbitrary state. This evaluation inherently includes both the forward and reverse processes of stacking blocks, determining whether an agent that can stack blocks is also capable of returning the stacked blocks to an intermediate state. In the Ant Maze environment, we use 32 different positions within the maze these positions allows for comprehensive assessment of the Ant Robot's world model learning for the maze structure. This evaluation not only measures its ability reach the final room but also its capability to return to previous rooms from intermediate positions. We evaluate 1024 combinations of starting and goal positions, conducting 10 evaluations for each combination, resulting an average success rate computed over 10,240 experiments. In the Walker environment, we use all evaluation goals (+:13, + 16, 19, +22 +:25, +:28) as a candidate set for starting and goal positions. This set can form total of 144 different combinations of starting and goal positions, providing a thorough assessment of the Walker robot's ability to move forward and backward, as well as its precision in position judgment. See Table for specific results of MUN and all baselines.
candidate for starting and goal positions. Evaluating navigation between
21

---

Table 3: Success rate of navigation experiments on 3-Block Stacking
Environment
Success rate
MUN
3-Block Stacking
95%
MUN-noDAD
3-Block Stacking
81% %
OC-Dreamer
3-Block Stacking
56%
MEGA-G
3-Block Stacking
42%
PEG-G
3-Block Stacking
47%
Table 4: Success rate of navigation experiments Ant Maze
Environment
Success rate
MUN
Ant-Maze
96%
MUN-noDAD
Ant-Maze
89%
GC-Dreamer
Ant-Maze
75%
MEGA-G
Ant-Maze
94%
PEG-G
Ant-Maze
93%
Table 5: Success rate of navigation experiments Walker
Environment
Success rate
MUN
Walker
89%
MUN-noDAD
Walker
73%
BC-Dreamer
Walker
67%
MEGA-G
Walker
81%
PEG-G
Walker
62%
We observe that MUN significantly outperforms other baselines in navigation experiments across all three environments, demonstrating exceptional contribution learning comprehensive world models and policies.
F.3
World Model Assessment
shows the single-step prediction error of learned world models. We randomly sample le4 transition tuples (s,, a, 81+1) within the replay buffers from all of our baselines (MUN, MUN- noDAD, GC-Dreamer, MEGA-G, and PEG-G) to form validation dataset. Table repports the mean squared error on this dataset.
Table shows the compounding error (multistep prediction error) of learned world models for evaluation when generating the same length simulated trajectories. More specifically, assume real trajectory of length h is denoted as (Su. ao $1, Sn). For a learned model M, we sample from 8o and generate forward rollouts (ša, ag $1, ..., 5n) where so so and for i < 0, 8 M ($o, a;).
Then the corresponding compounding error of is defined as Σ || s We set h to be the maximum number of timesteps in environments. We evaluated the compounding prediction error of the learned world models by generating 500 trajectories for each benchmark, simulated on both the models and the real environments.
In Tables 6 and 7 we used the final world models trained by all methods after the same number of environment interaction steps. These results provide a quantitative comparison of the world model prediction quality between MUN and the baselines across our benchmarks. The world models trained by MUN show much smaller generalization gap the real environment compared to goal-conditioned Dreamer (and the other baselines). Consequently, MUN can effectively leverage these world models to train control policies that generalize well to the real environment. This explains
22

---

the superior task success rates of MUN compared to the baselines in our experiment. Fig 10 also provides more information about the world model compound prediction error.
Wir Medell Pradiction
Wo Model Prodiction
0l0u-d 1443
1.73
PCD-0
MOOK 3
1.25
0.00
200p 124
140
200p
(a) Trajectories 3-Block Stacking
(b) Trajectorics Block Rotation
Figure 10: Fig(a) and Fig(b) illustrate the imagined and real environment trajectories for 3-Block Stacking and Block Rotation respectively, starting from the same initial state. Among the baselines, MUN demonstrates the smallest compound model error with respect to the ground truth trajectories. The X-axis represents the trajectory steps. In Fig(a), the Y-axis represents the sum of the heights of the three blocks. MUN's world model outperforms other methods in predicting the correct locations of the three blocks. In Fig(b), the Y-axis represents the position of the block in the coordinate. MUN's world model outperforms other methods in predicting the correct position of the block.
Table 6: One-step Model Prediction Error.
MUN
MUN-noDAD
PEG-G
MEGA-G
GC-Dreamer
Ant Maze
1.6740
1,9751
2.1154
2.2416
2.9666
Walker
0.8165
0,9971
1,4759
1,2353
2.1824
3-Block Stacking
0,0070
0.0071
0,0476
0.0853
0.0392
Rotate Block
1.0570
1.5609
1.7753
1.9433
2.3723
Rotate Pen
0.6708
1,1999
1.9622
2.8598
1,8359
Fetch Slide
0.0094
0.0108
0.0132
0.0164
0.0169
Table 7: Compound Model Prediction Error.
MUN
MUN-noDAD
PEG-G
MEGA-G
GC-Dreamer
Ant Maze
18.83
22.42
29.42
23.69
40.36
Walker
13.03
16.72
26.54
21.21
39.72
3-Block Stacking
0.45
0.55
0.70
0.95
0.94
Rotate Block
11.55
12.86
14.38
14.13
15.06
Rotate Pen
4.63
6.10
7.40
9.85
9.36
Fetch Slide
1.687
1.648
2.195
2.856
2.304
23

---

NeurIPS Paper Checklist
1. Claims
Question: Do the main claims made in the abstract and introduction accurately reflect the paper's contributions and scope?
Answer: [Yes]
Justification: Our abstract and introduction accurately represent the primary contributions of the paper, which include the development of the MUN algorithm designed to improve goal-conditioned reinforcement learning through enhanced world modeling and exploration capabilities. The introduction outlines the key challenges in GCRL, specifically with sparse rewards, and how MUN addresses these by facilitating effective state transitions between arbitrary subgoal states in the replay buffer. These claims are well-supported by the theoretical underpinnings and experimental results presented in the paper, reflecting the scope and impact of the proposed method.
Guidelines:
The answer NA means that the abstract and introduction do not include the claims made in the paper.
The abstract and/or introduction should clearly state the claims made, including the contributions made in the paper and important assumptions and limitations. A No or
NA answer to this question will not perceived well by the reviewers.
The claims made should match theoretical and experimental results, and reflect how much the results can be expected to generalize to other settings.
It is fine to include aspirational goals as motivation long it is clear that these goals are not attained by the paper.
2. Limitations
Question:
Does the paper discuss the limitations of the work performed by the authors?
Answer: [Yes]
Justification: Our paper thoroughly discusses the limitations of the MUN framework in Appendix. We highlight the dependency on an efficient strategy for discovering key states, pointing out that while DAD is effective, it will also identify irrelevant states in tasks with complex action spaces or weak correlations between goal space and action space. Addi- tionally, we mention the potential for applying MUN to model-free reinforcement learning methods, which do not require learning a world model and have simpler architectures.
Guidelines:
The answer NA means that the paper has limitation while the answer No means that the paper has limitations, but those are not discussed in the paper.
The authors are encouraged to create separate "Limitations" section in their paper.
The paper should point out any strong assumptions and how robust the results are to violations of these assumptions (e.g., independence assumptions, noiseless settings, model well-specification, asymptotic approximations only holding locally). The authors should reflect on how these assumptions might be violated in practice and what the implications would be.
The authors should reflect on the scope of the claims made, e.g., if the approach was only tested on few datasets or with few runs. In general, empirical results often depend on implicit assumptions, which should be articulated.
The authors should reflect on the factors that influence the performance of the approach. For example, facial recognition algorithm may perform poorly when image resolution low or images are taken in low lighting. Or speech-to-texx system might not be used reliably to provide closed captions for online lectures because it fails handle technical jargon.
The authors should discuss the computational efficiency of the proposed algorithms and how they scale with dataset size.
If applicable, the authors should discuss possible limitations of their approach to address problems of privacy and faimess.
24

---

While the authors might fear that complete honesty about limitations might be used by reviewers as grounds for rejection, worse outcome might be that reviewers discover limitations that aren't acknowledged in the paper. The authors should use their best judgment and recognize that individual actions in favor of transparency play an impor- tant role in developing norms that preserve the integrity of the community. Reviewers will be specifically instructed to not penalize honesty concerning limitations.
3. Theory Assumptions and Proofs
Question: For each theoretical result, does the paper provide the full set of assumptions and a complete (and correct) proof?
Answer: [NA]
Justification: Our paper does not include theoretical results.
Guidelines: In the Experiment section and the Appendix of our paper, we provide a detailed description of our experimental procedures and configurations. This includes all sources and modifications of the test environments, pseudocode and implementation methods for all baselines, the equipment and memory used, as well as the specific values of the required hy pypeparraeters. Additionally, we have open-sourced our code, which can be found in the Reproducibility Statement section.
The answer NA means that the paper does not include theoretical results.
All the theorems, formulas, and proofs in the paper should be numbered and cross- referenced.
All assumptions should be clearly stated or referenced in the statement of any theorems.
The proofs can either appear in the main paper or the supplemental material, but if they appear in the supplemental material, the authors are encouraged to provide short proof sketch to provide intuition.
Inversely, any informal proof provided in the core of the paper should be complemented by formal proofs provided in appendix or supplemental material.
Theorems and Lemmas that the proof relies upon should be properly referenced.
4. Experimental Result Reproducibility
Does the paper fully disclose all the information needed to reproduce the main ex- perimental results of the paper to the extent that it affects the main claims and/or conclusions of the paper (regardless of whether the code and data are provided or not)?
Question:
Answer: [Yes]
Justification: In the Experiment section and appendix of our paper, we elaborate on the procedure and configuration of our experiments. This includes the sources and modifications of all testing environments, pseudo code and implementation methods for all baselines, the devices and memory utilized, as well as specific values of hyperparameters employed. Concurrently, we have open-sourced our code; please refer to the Reproducibility Statement section for further details.
Guidelines:
The answer NA means that the paper does not include experiments.
If the paper includes experiments, a No answer to this question will not be perceived well by the reviewers: Making the paper reproducible is important, regardless of whether the code and data are provided or not.
If the contribution is dataset and/or model, the authors should describe the steps taken to make their results reproducible or verifiable.
Depending on the contribution, reproducibility can be accomplished in various ways. For example, if the contribution is a novel architecture, describing the architecture fully might suffice, or if the contribution is specific model and empirical evaluation, it may be necessary to either make possible for others to replicate the model with the same dataset, or provide access to the model. In general. releasing code and data is often one good way to accomplish this, but reproducibility can also be provided via detailed instructions for how to replicate the results, access to a hosted model (e.g., in the case of large language model), releasing of a model checkpoint, or other means that are appropriate to the research performed.
25

---

reproducible, they While NeurIPS does not require releasing code, the conference does require all submis- sions to provide some reasonable avenue for reproducibility, which may depend on the nature of the contribution. For example
(a) If the contribution is primarily a new algorithm, the paper should make it clear how to reproduce that algorithm.
(b) If the contribution is primarily
model architecture, the paper should describe the architecture clearly and fully.
(c) the contribution is a new model (c.g., a large language model), then there should either be way to access this model for reproducing the results or way to reproduce the model (e.g., with an open-source dataset or instructions for how to construct the dataset).
(d) We recognize that reproducibility may be tricky in some cases, in which case authors are welcome deseribe the particular way they provide for reproducibility, In the case of closed-source models, it may be that access to the model is limited in some way (e.g., to registered users), but it should be possible for other researchers to have some path to reproducing or verifying the results.
5. Open access to data and code
Question: Does the paper provide open access to the data and code, with sufficient instruc- tions to faithfully reproduce the main experimental results, as described in supplemental material?
Answer: [Yes]
Justification: As we mentioned the previous justification, we have not only open-sourced our code but also provided detailed steps and settings for reproducing our main experimental results. In the Experiment section and Appendix, we elaborate on the sources and modifi- cations of the environments, baseline implementation details, and MUN implementation specifics.
Guidelines:
The answer NA means that paper does not include experiments requiring code.
Please see the NeurlPS code and data submission guidelines (https://nips.co cc/
for more details.
While we encourage the release of code and data, we understand that this might not be possible, "No" is an acceptable answer. Papers cannot be rejected simply for not including code, unless this is central to the contribution (e.g., for new open-source benchmark).
The instructions should contain the exact command and environment needed to run to reproduce the results. See the NeurIPS code and data submission guidelines (https
/ //nips.co httpssccc/ppo wttccc/ppuuu ppbulcc/guudeeee
SooodeSubmissioPPono
for more details.
The authors should provide instructions data access and preparation, including how to access the raw data, preprocessed data, intermediate data, and generated data, etc.
The authors should provide scripts to reproduce all experimental results for the new proposed method and baselines. If only subset of experiments should state which ones are omitted from the script and why.
At submission time, to preserve anonymity, the authors should release anonymized versions (if applicable).
Providing as much information as possible in supplemental material (appended the paper) is recommended, but including URLs to data and code is permitted.
6. Experimental Setting/Details
Question: Does the paper specify all the training and test details (c.g., data splits, hyper- parameters, how they were chosen, type of optimizer, etc.) necessary to understand the results?
Answer: [Yes]
Justification: We provide comprehensive details regarding the hyperparameters essential for understanding the experiments, including those specific to our MUN framework. The table presented (Fig outlines these hyperparameters for each task, facilitating reproducibility and comparison.
26

---

Guidelines:
The answer NA means that the paper does not include experiments.
The experimental setting should be presented in the core of the paper to a level of detail that is necessary appreciate the results and make sense of them.
The full details can be provided either with the code, in appendix, as supplemental material.
7. Experiment Statistical Significance
Question: Does the paper report error bars suitably and correctly defined or other appropriate information about the statistical significance of the experiments?
Answer: [Yes]
Justification: We conducted each experiment minimum of five times using different random seeds, and upon plotting the results, as demonstrated in the Experiment section, we incorporated the experimental error. The solid line denotes the average success rate, while the shaded region signifies the standard deviation among the repeated experimental outcomes.
Guidelines:
The answer NA means that the paper does not include experiments.
The authors should answer "Yes" if the results are accompanied by error bars, confi- dence intervals, or statistical significance tests, least for the experiments that support the main claims of the paper.
The factors of variability that the error bars are capturing should be clearly stated (for example, train/test split, initialization, random drawing of some parameter, or overall run with given experimental conditions).
The method for calculating the error bars should be explained (closed form formula, call to a library function, bootstrap, etc.)
The assumptions made should be given (e.g., Normally distributed errors).
It should be clear whether the error bar is the standard deviation or the standard error of the mean.
is OK to report 1-sigma error bars, but one should state it. The authors should preferably report a 2-sigma error bar than state that they have 96% CI, if the hypothesis of Normality errors is not verified.
For asymmetric distributions, the authors should be careful not to show in tables or figures symmetric error bars that would yield results that out of range (e.g. negative error rates).
If error bars are reported in tables or plots, The authors should explain in the text how they were calculated and reference the corresponding figures tables in the text.
8. Experiments Compute Resources
For each experiment, does the paper provide sufficient information on the com- puter resources (type of compute workers, memory, time of execution) needed reproduce the experiments?
Question:
Answer: [Yes]
Justification: We clearly specifies the computer resources (Nvidia A100 GPU) and the amount of GPU memory required (approximately 3GB). Additionally, we provides detailed information the runtime of each experiment in Appendix.
Guidelines:
The answer NA means that the paper does not include experiments.
The paper should indicate the type of compute workers CPU or GPU, internal cluster, cloud provider, including relevant memory and storage.
The paper should provide the amount of compute required for each of the individual experimental runs as well as estimate the total compute.
The paper should disclose whether the full research project required more compute than the experiments reported the paper (e.g., preliminary or failed experiments that didn't 't make into the paper).
27

---

9. Code Of Ethics
Question: Does the research conducted in the paper conform, in every respect, with the NeurlPS Code of Ethics httppp////oourpsscccc////icciiie
Answer: [Yes]
The research conducted in our paper aligns with the NeurlPS Code of Ethics. We have thoroughly reviewed the guidelines and ensured that our research adheres to ethical standards. Additionally, we have implemented measures to safeguard anonymity and comply with pertinent laws and regulations.
Justification:
Guidelines:
The answer NA means that the authors have not reviewed the NeurIPS Code of Ethics.
If the authors answer No, they should explain the special cireumstances that require deviation from the Code of Ethics.
The authors should make sure to preserve anonymity (e.g., if there is special consid- eration due to laws or regulations their jurisdiction).
10. Broader Impacts
Question: Does the paper discuss both potential positive societal impacts and negative societal impacts of the work performed?
Answer: [NA]
Justification: Our research aims to address the exploration problem in Reinforcement Learning (RL) within the GCRL environment. It is currently in the theoretical research stage and has minimal societal impact.
Guidelines:
The answer NA means that there is no societal impact of the work performed.
If the authors answer NA or No, they should explain why their work has no societal impact or why the paper does not address societal impact.
Examples of negative societal impacts include potential malicious or unintended uses (e.g., disinformation, generating fake profiles, surveillance), fairness considerations (e.g., deployment of technologies that could make decisions that unfairly impact specific groups), privacy considerations, and security considerations.
The conference expects that many papers will be foundational research and tied to particular applications, let alone deployments. However, if there is direct path to any negative applications, the authors should point it out. For example, it is legitimate to point out that improvement in the quality of generative models could be used to generate deepfakes for disinformation. On the other hand, it is not needed to point out that generic algorithm for optimizing neural networks could enable people to train models that generate Deepfakes faster.
The authors should consider possible harms that could arise when the technology is being used intended and functioning correctly, harms that could arise when the technology is being used as intended but gives incorrect results, and harms following from (intentional or unintentional) misuse of the technology.
there are negative societal impacts, the authors could also discuss possible mitigation strategies (e.g., gated release of models, providing defenses in addition to attacks, mechanisms for monitoring misuse, mechanisms to monitor how a system learns from feedback over time, improving the efficiency and accessibility of ML).
11. Safeguards
Does the paper deseribe safeguards that have been put in place for responsible release of data models that have high risk for misuse (e.g., pretrained language models, image generators, or scraped datasets)?
Question:
Answer: [NA]
Justification: Our paper poses no such risks.
Guidelines:
The answer NA means that the paper poses such risks.
28

---

Released models that have high risk for misuse or dual-use use should be released with necessary safeguards to allow for controlled use of the model. for example by requiring that users adhere to usage guidelines or restrictions to access the model or implementing safety filters.
Datasets that have been scraped from the Internet could pose safety risks. The authors should describe how they avoided releasing unsafe images.
We recognize that providing effective safeguards challenging, and many papers do not require this, but encourage authors to take this into account and make best faith effort.
12. Licenses for existing assets
Question: Are the creators or original owners of assets (e.g.. code, data, models), used in the paper. properly credited and are the license and terms of use explicitly mentioned and properly respected?
Answer: [Yes]
Our paper properly credits the creators original owners of assets used, including code, data, and models. The licenses and terms of use are explicitly respected. Specifically, cite the original papers for code packages datasets used, state the version the assets, and include URLs where possible.
Justification:
Guidelines:
The answer NA means that the paper does not use existing assets.
The authors should cite the original paper that produced the code package dataset.
The authors should state which version asset used and, possible, include URL.
The name of the license (e.g., CC-BY 4.0) should be included for each asset.
For scraped data from a particular source (e.g.. website). the copyright and terms of service of that source should be provided.
assets are released, the license, copyright information, and terms use in package should be provided. For popular datasets, paperswi thcode. com datasets has curated licenses for some datasets. Their licensing guide can help determine the license of dataset.
For existing datasets that are re- re-packaged, both the original license and the license of the derived asset (if it has changed) should be provided.
If this information is available online, the authors encouraged to reach out the asset's 's creators.
13. New Assets
Question: Are new assets introduced in the paper well documented and is the documentation provided alongside the assets?
Answer: [Yes]
Justification: We have documented our code and provided detailed instructions on its usage, licenses, and permissible scope of use. Additionally, have included the documentation alongside the assets to ensure accessibility and clarity for users.
Guidelines:
The answer NA means that the paper does not release new assets.
Researchers should communicate the details of the dataset/code/model part of their submissions via structured templates. This includes details about training, license, limitations, etc.
The paper should discuss whether and how consent was obtained from people whose asset used.
At submission time, remember anonymize your assets (if applicable). You can either create an anonymized URL or include an anonymized zip file.
14. Crowdsourcing and Research with Human Subjects

---

Question: For crowdsourcing experiments and research with human subjects, does the paper include the full text of instructions given to participants and screenshots, if applicable, as well as details about compensation (if any)?
Answer: [NA]
Justification: Our paper not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Including this information in the supplemental material is fine, but if the main contribu- tion of the paper involves human subjects, then as much detail as possible should be included in the main paper.
According to the NeurIPS Code of Ethies, workers involved in data collection, curation, or other labor should be paid at least the minimum wage in the country of the data collector.
15. Institutional Review Board (IRB) Approvals or Equivalent for Research with Human Subjects
Does the paper describe potential risks incurred by study participants, whether such risks were disclosed to the subjects, and whether Institutional Review Board (IRB) approvals (or equivalent approval/review based the requirements of your country or institution) were obtained?
Question:
Answer: [NA]
Justification: Our paper does not involve crowdsourcing nor research with human subjects.
Guidelines:
The answer NA means that the paper does not involve crowdsourcing research with human subjects.
Depending on the country in which research is conducted, IRB approval (or equivalent) may be required for any human subjects research. If you obtained IRB approval, you should clearly state this in the paper.
We recognize that the procedures for this may vary significantly between institutions and locations, and we expect authors to adhere to the NeurIPS Code of Ethics and the guidelines for their institution.
For initial submissions, do not include any information that would break anonymity (if applicable), such as the institution conducting the review.
30


Reviews:
[{"review_id": "wcG2zVkIml", "reviewer": "Reviewer_Rfj2", "summary": "**After rebuttal the score was changed from reject to weak accept.**\n\nThe paper proposes to use the experiences stored in an RL replay buffer differently when training a world model. The change is to attempt to not only use the experience in a \"forward\" direction but also a \"backward\" and \"across traces\" manner. The hope is that this would result in a world model capable of better predicting the behavior of the real world when generalizing to novel trajectories. A secondary addition is a heuristic to select exploration goals which relies on the assumption that large state space differences correspond to meaningful subgoals.", "questions": "- How does the proposed subgoal heuristic compare to more complex approaches such as the ones mentioned in the paper?\n- Would it be possible to automatically detect whether the proposed backwards / sideways trace playback for world model learning is applicable for the given problem, such that the approach can be used in scenarios where its assumptions are met", "limitations": "The paper does not contain a limitations section in the body of the paper, and the limitation section in the appendix does not truly address the assumptions and limitations stemming from them outlined in this review.", "rating": 6, "confidence": 4, "soundness": 2, "presentation": 2, "contribution": 2, "strengths": "- Improving the generalization capability of the world model is a sensible approach", "weaknesses": "- Strong implied assumptions about symmetry of the problem setup not considered or discussed.\n- Assumptions of the proposed solutions not discussed.\n- Detail of the method remain at a very high level\n\nThe core idea of the paper is to improve the world model by using replay buffer trajectories backwards or transitioning between states on different trajectories. If these transitions are correct they will improve the quality of the world model as it has more data diversity. However, the paper provides no discussion or insight into how these transitions can be verified to be correct. The proposed approach can work, however, only for specific kinds of problem setups where the transitions are symmetric. For tasks such as cube stacking with pose-based states this idea is reasonable. For many other tasks that include real physical dynamics, such as driving, pouring liquids, or non-prehensile manipulation it is unlikely that this approach would produce good results. This limitation has to be clearly addressed and explored in the paper. It doesn't mean that the idea is invalid because of it, but without a thorough discussion of it the paper is incomplete.\n\nThe subgoal discovery heuristic similarly makes the assumption that distant points in the state space correspond to crucial key states for a task. While the experiments show that it can be beneficial, there is no theoretical or intuitive evidence provided that this should work in general. Whether this heuristic is beneficial will heavily depend on the design or embedding of the state space. While one can follow the basic idea behind the heuristic more rigorous discussion and theoretical work is required on this front.\n\nWhile the high-level ideas of the paper are easy to understand and follow the detailed information is lacking or not clear enough. While the algorithms are ok some of the variables used within, especially Algorithm 1, are not described and it is not entirely clear which steps are done in the world model and which in the real/simulated environment.\n\nThe experiments show how the proposed approach can outperform the selected baselines in certain tasks. Despite this there are several aspects that I would have expected the experiments to validate. The first is an evaluation of the world model. The main idea and contribution of the paper is a mechanism to build a better world model. Yet that quality of the world model is never properly evaluated. Figure 6 attempts this to some extent on the task most amenable of all to the implied assumptions. As such a thorough evaluation of the quality of the world model itself would be needed as part of the discussion and evaluation of the method's assumptions and limitations. Another is the evaluation of the exploration goal selection mechanism. While the experiments compare with methods utilizing approaches in the spirit of go-explore, there is no ablation of the proposed system using that mechanism. As such it is impossible to assess what, if any, benefit that strategy has as the results are confounded with the modified world model. An interesting aspect the plots in Figure 4 do not show is the performance achievable by different methods upon convergence. While faster learning is always good it is also interesting to know what performance difference can be achieved at convergence. The evaluation of the subgoal discovery policy is unconvincing. While the images show images of interesting steps it is unclear that they are meaningful for the policy learning. A rigorous quantitative evaluation of this would be needed to gain insights into the heuristic.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**R1. The subgoal discovery heuristic assumes that distant points in the state space correspond to crucial key states. Whether this heuristic is beneficial will depend on the design or embedding of the state space**\n\nMUN does *not* directly find key subgoal states from the state space and, therefore, does not depend on the embedding of the state space for key subgoal discovery. The DAD strategy in MUN discovers actions that significantly differ along trajectories in the replay buffer and extracts the corresponding states triggering these actions as potential key subgoal states. Our strategy is based on the observation that the agent often performs distinct actions to trigger key subgoal states e.g. closing grippers to grasp objects. See R2 for more examples.\n\n**R2. Strong assumptions about symmetry of the problem setup not considered. For many other tasks that include real physical dynamics, such as driving, pouring liquids, or non-prehensile manipulation it is unlikely that this approach would produce good results**\n\nMUN is effective for non-prehensile manipulation tasks. To demonstrate this, we applied MUN to the Gymnasium Fetch Slide task, where a manipulator slides a puck to a target on a slippery table. This environment has asymmetric state transitions: when the puck is slid outside the robot's workspace, the manipulator cannot reach the puck's position to slide it backward due to physical constraints. The training performance in terms of task success rates is depicted in Fig. 1(g) in the global response. MUN significantly outperforms the goal-conditioned Dreamer baseline. By examination, we found MUN with the DAD strategy can discover key subgoals for this task like contacting the puck, placing the manipulator at different angles around the puck, and stabilizing the manipulator upon reaching the goal (these key states result from distinct actions). MUN enables learning state transitions between these key subgoals to discover the high-level task structure. It learns a generalizable world model that handles sliding the puck between any positions within the workspace and predicts low probabilities for infeasible transitions from puck positions outside the workspace. Particularly, it enables the agent to hit the puck multiple times if it is within its workspace, thereby improving task success rates. That said, the current goal selection mechanism in MUN lacks a mechanism to filter out infeasible goals from a current state, potentially affecting sample efficiency. Please refer to R3 as a possible solution. We will clarify this limitation in the revised paper.\n\n**R3. Would it be possible to automatically detect whether the proposed backwards / sideways trace playback for world model learning is applicable for the given problem?**\n\nYes, we can track the learning progress of a goal state $g$ from a state $s$ by monitoring changes in the expected discounted returns (value functions) in the learned world models across training episodes. We could sample goal states based on their learning progress in the goal sampling step at Line 10 of Algorithm 1. If progress towards $g$ stagnates, due to factors like asymmetric state transitions, we can reduce its selection probability. We plan to explore this approach in future work.\n\n**R4. Some of the variables used within Algorithm 1 are not described and it is not clear which steps are done in the world model and which in the real environment.**\n\nLine 19 in Algorithm 1 is the only step we perform simulation in a world model for policy training. The other steps occur in the real environment. We will improve the clarity of Algorithm 1.\n\n**R5. Evaluation of the world model**\n\nPlease see Tables 1 and 2 in the global response.\n\n**R6. Figure 4 does not show the performance upon convergence**\n\nWe have extended the training steps. Fig.1 (a)-(f) in the global response shows that MUN outperforms the baselines in both final performance and sample efficiency.\n\n**R7. No ablation of the exploration goal selection mechanism**\n\nPlease see the additional ablations (MUN-KeyObs and MUN-Ns-3) in the global rebuttal.\n\n**R8. How does the proposed subgoal heuristic compare to more complex approaches such as the ones mentioned in the paper?**\n\nThe prior work (Paul et al. 2019) learns subgoals from expert trajectories, which we do not assume are available. Please refer to our global response (the MUN-KeyObs ablation) for the quantitative comparison with (Zhang et al. 2021) which directly discovers key subgoals scattered across the state space.\n\n**R9. The evaluation of subgoal discovery is unconvincing. It is unclear whether the subgoals are meaningful for policy learning.**\n\nTo evaluate DAD (the subgoal discovery Algorithm 2), we designed the MUN-noDAD ablation to assess how the subgoals improve policy training. MUN-noDAD replaces DAD with a simpler method that selects goal states at fixed time intervals in trajectories from the replay buffer. As shown in Fig. 4, MUN outperforms MUN-noDAD across all benchmarks, especially in the challenging 3-block stacking and high-dimensional block rotation environments. Tables 1 and 2 in the global response illustrate that the world models trained by MUN have a smaller generalization gap to real environments compared to those trained by MUN-noDAD. Consequently, MUN's superior world models enable more effective training of control policies that generalize better to real-world scenarios. We further evaluated MUN and MUN-noDAD on tasks not encountered during training to assess the model generalizability, as discussed in Sections 4.5 and F.5 of the appendix. These novel tasks feature different initial state and final goal configurations compared to the training tasks. MUN achieved an average success rate improvement of 16.5% over MUN-noDAD on these unseen tasks, as shown in Tables 3, 4, and 5 in the appendix. We will revise the paper to clarify the comparison with MUN-noDAD.", "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": "HPEzKwlk6s", "reviewer": "Reviewer_BNZB", "summary": "This paper introduces a novel goal-directed exploration algorithm called MUN to address the challenges of efficient exploration in long-horizon, sparse-reward environments within the context of goal-conditioned reinforcement learning (GCRL). The key insight is that improving the generalizability of learned world models can significantly enhance the agent's capacity to explore and navigate in the real environment. MUN focuses on training world models that can accurately characterize state transitions between arbitrary subgoal states, whether by retracing along recorded trajectories or transitioning between states on separate trajectories. Additionally, the paper presents a practical strategy called DAD (Distinct Action Discovery) to identify pivotal subgoal states that represent critical milestones in completing complex tasks. By training world models and policies for seamless transitions between these key subgoals, MUN is able to learn generalizable goal-conditioned policies that can adapt to novel test scenarios.", "questions": "Questions and concerns are stated above.", "limitations": "Yes.", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 2, "strengths": "1. This paper is well-written with thorough experimental analysis. \n\n2. The DAD method for identifying pivotal subgoal states is a practical and effective strategy for enhancing policy generalization.", "weaknesses": "1. MUN does not seem to outperform the baselines significantly as seen in Figure 5. \n\n2. DAD as an acronym conflicts with another work in the same literature, \"Dynamic-Aware Discovery of Skills\" (Sharma et al., 2019) \n\n3. The distinction from Go-Explore approach is not super clear. An pictorial explanation may help a lot. The \"Comparison with Go-Explore\" paragraph is a bit dense. Given that N_s is always set to 2 in the paper, it's not clear how big of a distinction this makes. \n\n4. Ablations are lacking besides MUN-noDAD; there are many moving pieces in the algorithm, and it would be nice to see their effects on the final method performance. \n\n5. The paper uses the phrasing \"as the quality of the world model improves\" many times, but it's not clear what that means in this paper.; no prediction quality of the dynamics model was studied.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**R1. MUN does not seem to outperform the baselines significantly as seen in Figure 5 (Figure 4?).**\n\nFigure 4 might give a misleading impression as we blend the baselines and ablations in the same image. In this figure, the tool with the closest performance to MUN is our ablation, MUN-noDAD. We will revise the paper to separate the comparisons between the baselines and the ablation to more clearly demonstrate the performance advantages of MUN.\n\nThe final performance in terms of task success rates of MUN and the baselines (evaluated over 100 seeds) is summarized below:\n\n| Env/Method        | MUN | PEG-G | MEGA-G | GC-Dreamer |\n|-------------------|-----|-------|--------|------------|\n| **Ant Maze**      | 97% | 80%   | 87%    | 76%        |\n| **Point Maze**    | 100%| 90%   | 100%   | 22%        |\n| **Walker**        | 91% | 79%   | 83%    | 72%        |\n| **3-Block Stacking** | 98% | 61%   | 39%    | 70%        |\n| **Rotate Block**  | 86% | 70%   | 71%    | 62%        |\n| **Rotate Pen**    | 72% | 60%   | 45%    | 49%        |\n\nIt can be seen that MUN outperforms the baselines by a large margin in challenging environments such as 3-Block Stacking, Block Rotation, and Pen Rotation (with 61 state dimensions and 20 action dimensions).\n\n**R2. DAD as an acronym conflicts with another work in the same literature, \"Dynamic-Aware Discovery of Skills\" (Sharma et al., 2019)**\n\nThanks for pointing this out. We will rename DAD in the revised paper.\n\n**R3. The distinction from Go-Explore approach is not super clear.**\n\nA Go-Explore agent in training is directed to traverse to a goal state with high exploration potential, sampled from the replay buffer. Once at this goal state, the agent switches to a separate exploration policy to further explore the environment, which is trained to maximize an intrinsic exploration reward. In contrast, a MUN agent eliminates the need for a separate exploration policy and intrinsic exploration rewards. It is directed to sequentially traverse two goal states sampled from the replay buffer to explore the environment during each training episode. The trajectories collected from this exploration are used to train a world model, helping to close the generalization gap between the policy’s behavior in the real environment and the world model. This enables the MUN agent to effectively exploit the model for learning a policy that generalizes well to the real environment. We will visualize these key differences in the revised paper. \n\n**R4. Ablations are lacking besides MUN-noDAD; there are many moving pieces in the algorithm, and it would be nice to see their effects on the final method performance.**\n\nWe appreciate the reviewer for the great suggestion. We included additional ablation studies in the global response (the MUN-KeyObs and MUN-NS-3 ablations), which will be incorporated into the paper.\n\n**R5. The paper uses the phrasing \"as the quality of the world model improves\" many times, but it's not clear what that means in this paper.; no prediction quality of the dynamics model was studied.**\n\nThis is a great point! We included a comparison of one-step model prediction errors and compound model errors between MUN and the baselines in the global response Tables 1 and 2, which will be incorporated into the paper.", "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": "I3S7eKZW4R", "reviewer": "Reviewer_aZux", "summary": "The paper introduces MUN,  a novel goal-directed exploration algorithm for model-based reinforcement learning. MUN focuses on improving the quality of world models by learning to predict state transitions between any pair of states in the replay buffer. This is achieved by training the world model on a bidirectional replay buffer and identifying key subgoal states. By doing so, MUN enhances the generalizability of the learned world model and leads to more effective policy learning.\n\nKey Contributions:\n\nBidirectional World Model: Trains the world model to predict state transitions both forward and backward along trajectories, improving generalization.\n\nKey Subgoal Identification: Proposes a method called DAD to identify key subgoals, which are crucial for complex tasks.\n\nImproved Exploration: Demonstrates superior performance compared to baseline methods in various challenging environments.\n\nOverall, the paper presents a promising approach to improving the effectiveness of model-based reinforcement learning. Addressing the mentioned limitations and exploring potential extensions could further enhance the contribution of this work.", "questions": "How does the performance of MUN scale with the complexity of the environment (e.g., larger state and action spaces)?\n\nHave the authors explored different methods for identifying key subgoals, and how do they compare to DAD?\n\nWhat is the computational overhead of training a bidirectional world model compared to a traditional unidirectional model?\n\nHow does the choice of the world model architecture (RSSM) impact the performance of MUN?\n\nCan the authors provide a more detailed analysis of the ablation study comparing MUN with MUN-noDAD?", "limitations": "The paper has a section B of Appendix to handle this. No -ve societal impact for this work.", "rating": 8, "confidence": 5, "soundness": 4, "presentation": 4, "contribution": 3, "strengths": "Novel Approach: Addresses the limitations of traditional world models by focusing on bidirectional learning and key subgoal identification.\n\nStrong Empirical Results: Shows significant improvement over baseline methods in various complex environments.\n\nClear Explanation: Provides a detailed description of the proposed method and its components.\n\nThe codebase is shared and supplementary provides detailed report on experiments and environments.", "weaknesses": "Limited Baselines: While the paper compares to several baselines, a more comprehensive comparison with other state-of-the-art model-based RL methods could strengthen the claims. Real world environments in indoors and outdoors could have been explored.\n\nHyperparameter Sensitivity: The paper does not discuss the sensitivity of the method to hyperparameter tuning.\n\nLack of Theoretical Analysis: There is no theoretical analysis of the proposed method's convergence or optimality guarantees.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**R1. Limited Baselines**\n\nWe appreciate the reviewer’s suggestion. Our world model training strategy is compatible with any modern model-based RL framework. In this paper, we integrate it with the Dreamer framework due to Dreamer’s efficiency in optimizing policies directly in the latent space. We plan to explore integration with other model-based RL frameworks, such as MBPO and PlaNet, in a future extension of this work.\n\n**R2. Hyperparameter Sensitivity**\n\nWe implement MUN on top of the authors’ code from DreamerV2 (Hafner et al., 2021) and LEXA (Mendonca et al., 2021). We used the default hyperparameters for training the world model, policies, value functions, and temporal reward functions. Besides, MUN has three key hyperparameters: (1) the number of key subgoals ($N_{subgoals}$) to sample from the replay buffer using the DAD algorithm, (2) the number of subgoals ($N_s$) to sequentially explore during each training episode, and (3) the maximum number of timesteps ($T_s$) allocated for navigating to a specific subgoal. MUN is not sensitive to $N_{subgoals}$, and we set it to 20 for the most challenging benchmarks, such as 3-Block Stacking, Block Rotation, and Pen Rotation, to provide sufficient candidate subgoals. For $N_s$, we found that setting it higher than 3 slows down the training performance compared to our default value $N_s = 2$; see the global response (the MUN-NS-3 ablation) for quantitative comparisons. For $T_s$, we simply set it as the maximum episode length divided by $N_s$. More details can be found in Table 2 of the appendix.\n\n**R3. What is the computational overhead of training a bidirectional world model compared to a traditional unidirectional model?**\n\nMUN-noDAD periodically directs the agent to reach goal states sampled either from the replay buffer or from the environment's goal distribution, while standard model-based RL typically samples only from the latter. This goal state setting difference in MUN does not add computational overhead compared to standard model-based RL.\n\nMUN introduces additional overhead compared to MUN-noDAD, as the DAD algorithm employs the Farthest Point Sampling (FPS) technique for identifying key subgoals from a batch of samples in the replay buffer. FPS is computationally efficient, with a time complexity of $O(nk)$, where $n$ is the batch size and $k$ is the number of key subgoals to discover (a constant).\n\n**R4. How does the performance of MUN scale with the complexity of the environment (e.g., larger state and action spaces)?**\n\nAs mentioned in R3, MUN does not add significant overhead to standard model-based RL baselines such as Dreamer. MUN scales in accordance with Dreamer's handling of environmental complexity. The FPS algorithm, used to discover key subgoal states in DAD, operates in action spaces. Since the dimensionality of action spaces is typically smaller than that of state spaces, DAD remains scalable to environments with high-dimensional state spaces.\n\n**R5. Have the authors explored different methods for identifying key subgoals, and how do they compare to DAD?**\n\nWe considered applying DAD to discover key subgoals in the state spaces instead of the action spaces. Please refer to our global response (the MUN-KeyObs ablation) for details.\n\n**R6. How does the choice of the world model architecture (RSSM) impact the performance of MUN?**\n\nMUN improves the quality of a world model through its enriched replay buffers and a novel goal-sampling strategy. It is not tied to a specific world model architecture and can be applied to any model-based RL framework. We will clarify this in the paper.\n\n**R7. Can the authors provide a more detailed analysis of the ablation study comparing MUN with MUN-noDAD?**\n\nThe ablation MUN-noDAD replaces the goal sampling strategy DAD (Algorithm 2) with a simpler method that chooses goal states at fixed time intervals in trajectories sampled from the replay buffer. This ablation investigates the importance of identifying key subgoal states. It seeks to determine whether training world models from state transitions between these key subgoal states in MUN is essential, or if using any states from the replay buffer would suffice. As shown in Fig. 4, MUN outperforms MUN-noDAD on all of our benchmarks, particularly in the challenging 3-block stacking and high-dimensional block rotation environments. Fig. 7 (page 20) depicts the subgoals identified by the DAD algorithm during the training process for these benchmarks. For 3-block stacking, DAD identifies critical subgoals including block grasping, lifting, horizontal movement, vertical movement, and gripper release. For block rotation, DAD identifies crucial subgoals for finger movements that enable rotation. These results demonstrate that DAD can discover the high-level task structure, which is important for world modeling. Tables 1 and 2 in the global response quantitatively measure the world model one-step and compound errors of MUN and MUN-noDAD. The world models trained by MUN exhibit a smaller generalization gap to the real environment compared to those trained by MUN-noDAD. Consequently, MUN can effectively utilize these higher-quality world models to train control policies that generalize better to the real environment. We evaluated MUN and MUN-noDAD on tasks not encountered during training to assess the model generalizability, as discussed in Sections 4.5 and F.5 of the appendix. These novel tasks feature different initial state and final goal configurations compared to the training tasks. MUN achieved an average success rate improvement of 16.5% over MUN-noDAD on these unseen tasks, as shown in Tables 3, 4, and 5 in the appendix. We will incorporate this discussion into the paper.", "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": "D9t4Bj1ntn", "reviewer": "Reviewer_24DF", "summary": "The authors aim to improve learning in environments with sparse rewards by leveraging world models; this is achieved by proposing a novel exploration algorithm that as a result allows to create a richer buffer of experiences that are more appropriate for learning an accurate world model.", "questions": "* in figure 2 and section 3.1 the advantages of the bidirectional replay buffer are described. From the figure it is possible to see that with a bidirectional buffer there is a much wider range of possible trajectories covered. However, if an optimal trajectory can be found with a normal buffer, as illustrated in the figure, is there a need to learn a more complex space if the previous is enough to solve the task optimally?\n* how reasonable it is to find subgoals based on the FPS algorithm? just because there is a lot of variations in actions does it mean that these states correspond to different subgoals? \n* how do the authors define what is the resonable amount of subgoals over an episode?\t\n* finding subgoals is a popular topic in reinforcement learning; it is stated that the current approach is on the fact that key actions are needed at certain stages; is this easily generalizable? it sounds that this can depend a lot on the environment and might not very accurate in multiple cases", "limitations": "N/A", "rating": 5, "confidence": 2, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The paper is well organized and generally easy to follow. Tackling environments with sparse rewards is also important in RL.", "weaknesses": "* in line 123 it is stated that the representation might be inaccurate when moving backwards; it would be useful to ellaborate on why moving backwards needs to be taken into account, and not only forward\n* for completeness and better interpretability of the results, it would be interesting to see the world models that result as representations of the different states\n\n\nSome minor:\n* figure both refered to as Fig. and Figure (ex: line 117)\n* model-based RL both written as \"model-based RL\" (line 115) and \"Model-based RL\" (line 117)", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "We appreciate your insightful feedback and constructive comments!\n\n**R1. In line 123 it is stated that the representation might be inaccurate when moving backwards; it would be useful to ellaborate on why moving backwards needs to be taken into account, and not only forward**\n\nWe will revise the paper to improve the clarity. In model-based RL, a control policy is trained over *imagined rollouts* generated by a learned world model. A major issue is that the world model may generate hallucinated trajectories, leading to a significant discrepancy between the policy’s behavior in the model and in the real environment. Retraining the world model with new *forward* rollout samples from the real environment does not necessarily prevent these hallucinations. MUN addresses this by learning a world model that characterizes state transitions between any states within the replay buffer, thereby minimizing hallucinations if the replay buffer is representative. Tables 1 and 2 in the global response quantitatively measure the world model quality of MUN and GC-Dreamer, which only learns from forward rollout samples, across our benchmarks. The world models trained by MUN show a much smaller generalization gap to the real environment compared to GC-Dreamer. Consequently, MUN can effectively leverage these world models to train control policies that generalize well to the real environment.\n\n**R2. In figure 2 and section 3.1 the advantages of the bidirectional replay buffer are described. From the figure it is possible to see that with a bidirectional buffer there is a much wider range of possible trajectories covered. However, if an optimal trajectory can be found with a normal buffer, as illustrated in the figure, is there a need to learn a more complex space if the previous is enough to solve the task optimally?**\n\nIn model-based RL, the replay buffer (visualized in Figure 2) is not used directly to train a policy. Instead, it is used to train a world model, and the policy is trained over *imagined rollouts* generated by this model for sample efficiency. Consequently, having a few optimal trajectories in the replay buffer does not guarantee that the policy trained on the world model is optimal or generalizes well to the real environment. MUN enriches the replay buffer with state transitions by moving backward along recorded trajectories or across different trajectories. This approach trains high-quality world models that generalize well to the real environment (see R1 for more details) and uses these world models to learn optimal control policies.\n\n**R3. For completeness and better interpretability of the results, it would be interesting to see the world models that result as representations of the different states.**\n\nPlease refer to the PDF in the global rebuttal. Figures 1(h) and 1(i) illustrate both the imagined and real environment trajectories for 3-Block Stacking and Pen Rotation, starting from the same initial state. Among the baselines, MUN demonstrates the smallest compound model error. (See R1 and Tables 1 and 2 in the global rebuttal for more quantitative comparisons).\n\n**R4. How reasonable it is to find subgoals based on the FPS algorithm? just because there is a lot of variations in actions does it mean that these states correspond to different subgoals? Finding subgoals is a popular topic in reinforcement learning; it is stated that the current approach is on the fact that key actions are needed at certain stages; is this easily generalizable?**\n\nThe heuristic we use for subgoal generation is inspired by and most applicable to object manipulation and navigation tasks. We will clarify this limitation in the revised paper. In navigation tasks, changing the robot's direction inherently requires distinct actions compared to simply moving the robot in a single direction. Similarly, in manipulation tasks, actions such as moving the gripper to contact an object, pushing the object, and moving the gripper away upon goal reaching are all distinct. States resulting from the most distinct actions are likely subgoals, representing significant changes in events.\n\n**R5. How do the authors define what is the reasonable amount of subgoals over an episode?**\n\nWe define the number of subgoals $N_s$ within an episode to be 2. In our experience, setting $N_s$ greater than 2 does not improve training performance. Consider the task of block stacking with 3 potential subgoals found by DAD in Figure 5: (B) moving the grippers over a block, (C) closing the grippers to hold the block, and (D) moving the block to its goal region.  Learning how to transition from (B) to (C) and then from (C) to (D) by a policy is sufficient to solve the task. It is unnecessary to specifically learn transitioning from (B) to (C) to (D) in one episode. Based on our experience, setting $N_s$ to 3 slows down the training performance compared to our default value of $N_s = 2$. Please refer to the global response (the MUN-NS-3 ablation) for the quantitive comparisons. We will incorporate this discussion into the paper.", "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:
We greatly appreciate the valuable feedback and suggestions provided by the reviewers! We will begin by addressing the concerns raised by the majority of the reviewers in the global rebuttal. We will address the concerns of each reviewer in the individual review responses.

**1 Quantitive Measurement of World Model Predication Quality (Reviewer Rfj2, BNZB, 24DF)**

Table 1 shows the single-step prediction error of learned world models. We randomly sample state transition tuples $s_i, a_i, s\_{i+1}$ within the replay buffers from all of our baselines (MUN, MUN-noDAD, GC-Dreamer, MEGA-G, and PEG-G) to form a validation dataset. Table 1 reports the mean squared error on this dataset.

$\textbf{Table 1: One-step Model Predication Error}$
|                 | MUN    | MUN-noDAD | PEG-G  | MEGA-G | GC-Dreamer |
|-----------------|--------|-----------|--------|--------|------------|
| **Ant Maze**    | 1.6740 | 1.9751    | 2.1154 | 2.2416 | 2.9666     |
| **Point Maze**  | 0.0013 | 0.0013    | 0.0014 | 0.0011 | 0.0032     |
| **Walker**      | 0.8165 | 0.9971    | 1.4759 | 1.2353 | 2.1824     |
| **3-Block Stacking** | 0.0070 | 0.0071    | 0.0476 | 0.0853 | 0.0392     |
| **Rotate Block**| 1.0570 | 1.5609    | 1.7753 | 1.9433 | 2.3723     |
| **Rotate Pen**  | 0.6708 | 1.1999    | 1.9622 | 2.8598 | 1.8359     |

Table 2 shows the compounding error (multistep prediction error) of learned world models for evaluation when generating the same length simulated trajectories. More specifically, assume a real trajectory of length $h$ is denoted as $(s_0, a_0, s_1, \ldots, s_{h})$. For a learned model $M$, we sample from $s_0$ and generate forward rollouts $(\hat{s}_0, a_0, \hat{s}_1, \ldots, \hat{s}_h)$ where $\hat{s}_0 = s_0$ and for $i \geq 0$, $\hat{s}\_{i+1} = \mathcal{M} (\hat{s}_i, a_i)$. Then the corresponding compounding error of $\mathcal{M}$ is defined as $\frac{1}{h} \sum\_{i=1}^{h} \|\hat{s}_i - s_i\|_2^2$. We set $h$ to be the maximum number of timesteps in our environments.

$\textbf{Table 2: Compound Model Predication Error}$
|               | MUN  | MUN-noDAD | PEG-G | MEGA-G | GC-Dreamer |
|---------------|------|-----------|-------|--------|------------|
| **Ant Maze**       | 18.83 | 22.42      | 29.42 | 23.69   | 40.36      |
| **Point Maze**     | 5.59  | 5.43       | 6.32  | 4.74    | 9.57       |
| **Walker**         | 13.03 | 16.72      | 26.54 | 21.21   | 39.72      |
| **3-Block Stacking** | 0.45  | 0.55       | 0.70  | 0.95    | 0.94       |
| **Rotate Block**   | 11.55 | 12.86      | 14.38 | 14.13   | 15.06      |
| **Rotate Pen**     | 4.63  | 6.10       | 7.40  | 9.85    | 9.36       |

In Tables 1 and 2, we used the final world models trained by all methods after the same number of environment interaction steps. These results provide a quantitative comparison of the world model prediction quality between MUN and the baselines across our benchmarks. The world models trained by MUN show a much smaller generalization gap to the real environment compared to goal-conditioned Dreamer (and the other baselines). Consequently, MUN can effectively leverage these world models to train control policies that generalize well to the real environment. This explains the superior task success rates of MUN compared to the baselines in our experiment.

We will revise the paper to include this result and plot the model error curve throughout the training steps.

**2 Ablation Study for the Exploration Goal Selection Mechanism (Reviewer Rfj2, BNZB, aZux, 24DF)**

We conducted the following ablation studies to investigate MUN's exploration goal selection mechanism:

* I. Number of Goal States ($N_s$): MUN sequentially traverses $N_s = 2$ goal states sampled from the replay buffer to explore the environment during each training episode. We introduced an ablation **MUN-Ns-3** that sets $N_s=3$. The results in Fig.1 (d), (e), and (f) of the attached PDF show that setting $N_s$ greater than 2 slows down the training performance, supporting our claim in the paper it suffices to set $N_s = 2$.

* II. Key Subgoal Discovery: MUN discovers key subgoals for exploration as states in the replay buffer resulting from distinct actions in the action space. We provided an ablation **MUN-KeyObs** that discovers key subgoals directly from the state space as centroids of clusters of states in the replay buffer, following the strategy in [1]. The performance of this ablation as shown in Fig.1 (d), (e), and (f) of the attached PDF does not match MUN, highlighting that discovering key subgoals in the action space is both simpler and more effective.

We will incorporate these ablation studies into our paper.

[1] Zhang, L., Yang, G., and Stadie, B. World Model as a Graph: Learning Latent Landmarks for Planning. ICML 2021.

Meta-review: The paper introduces a goal-directed exploration algorithm for model-based reinforcement learning, designed to improve learning in sparse-reward, long-horizon environments. The approach learns world models by using experiences stored in the replay buffer not just in a forward direction but also backwards and across traces. This bidirectional training helps the world model better predict state transitions, even between arbitrary subgoal states. Another contribution is a heuristic that discovers novel actions and uses them to identify key subgoal states that are critical for completing complex tasks. By training on these subgoals, MUN improves exploration efficiency and leads to more effective goal-conditioned policies in novel scenarios.

Reviewers were generally positive about the approach, the novelty and the clarity of writing. However, there were initially some concerns about the paper, and the authors did a significant amount of work in the rebuttal phase. This included additional ablation studies, evaluating the quality of the learned model in isolation, as well as additional experiments on domains lacking the kinds of "symmetries" one might suspect the method requires (e.g. reversible dynamics). These results give further evidence to the strength and generality of the approach, although there are still open questions about under what conditions the method does (or does not) work. In particular, there seem to be many implicit assumptions that the conceptual ideas rely on, and without any theoretical foundation, it is difficult to understand when the method will succeed. 

Nonetheless, I think that the additional experiments conducted during the rebuttal sufficiently address the concerns of the reviewers and therefore recommend acceptance. Authors are strongly encouraged to incorporate these results into the final version of the paper, as well add additional explanations to improve clarity (particularly around Algorithm 1)
Decision: Accept (poster)