Learning World Models for Unconstrained Goal Navigation

We appreciate your insightful feedback and constructive comments!

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

We 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.

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?

In 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.

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.

Please 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).

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?

The 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.

R5. How do the authors define what is the reasonable amount of subgoals over an episode?

We 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.

We appreciate your insightful feedback and constructive comments!

R1. Limited Baselines

We 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.

R2. Hyperparameter Sensitivity

We 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.

R3. What is the computational overhead of training a bidirectional world model compared to a traditional unidirectional model?

MUN-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.

MUN 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).

R4. How does the performance of MUN scale with the complexity of the environment (e.g., larger state and action spaces)?

As 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.

R5. Have the authors explored different methods for identifying key subgoals, and how do they compare to DAD?

We 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.

R6. How does the choice of the world model architecture (RSSM) impact the performance of MUN?

MUN 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.

R7. Can the authors provide a more detailed analysis of the ablation study comparing MUN with MUN-noDAD?

The 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.

We appreciate your insightful feedback and constructive comments!

R1. MUN does not seem to outperform the baselines significantly as seen in Figure 5 (Figure 4?).

Figure 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.

The final performance in terms of task success rates of MUN and the baselines (evaluated over 100 seeds) is summarized below:

It 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).

R2. DAD as an acronym conflicts with another work in the same literature, "Dynamic-Aware Discovery of Skills" (Sharma et al., 2019)

Thanks for pointing this out. We will rename DAD in the revised paper.

R3. The distinction from Go-Explore approach is not super clear.

A 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.

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.

We 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.

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.

This 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.

We appreciate your insightful feedback and constructive comments!

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

MUN 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.

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

MUN 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.

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?

Yes, 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.

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.

Line 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.

R5. Evaluation of the world model

Please see Tables 1 and 2 in the global response.

R6. Figure 4 does not show the performance upon convergence

We 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.

R7. No ablation of the exploration goal selection mechanism

Please see the additional ablations (MUN-KeyObs and MUN-Ns-3) in the global rebuttal.

R8. How does the proposed subgoal heuristic compare to more complex approaches such as the ones mentioned in the paper?

The 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.

R9. The evaluation of subgoal discovery is unconvincing. It is unclear whether the subgoals are meaningful for policy learning.

To 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.

Thank you for considering our rebuttal and the newly presented experimental results. We will integrate the evaluation of the world model quality, the ablation study on key state selection, and the results from the new non-prehensile manipulation task into the paper. Additionally, we will strengthen the discussion on problem structures and the limitations of our approach as outlined in the rebuttals (e.g., emphasizing that MUN is primarily designed to address complex object manipulation and navigation tasks). We appreciate your advice for improving the quality of the paper.

In our rebuttal, we forgot to include the sizes of the validation datasets used to generate Tables 1 and 2.

For Table 1, we randomly sampled 10,000 state transition tuples $s_i, a_i, s_{i+1}$ from the replay buffers, ensuring an equal number of samples from each baseline (MUN, MUN-noDAD, GC-Dreamer, MEGA-G, and PEG-G), to form the validation dataset for measuring one-step model prediction errors.

For Table 2, we evaluated the compounding (multistep) prediction error of the learned world models by generating 500 trajectories for each benchmark, simulated on both the models and the real environments.