Efficiently learning to solve tasks in complex environments is a key challenge for reinforcement learning (RL) agents. We propose to decompose a complex environment using a task-agnostic world graphs, an abstraction that accelerates learning by enabling agents to focus exploration on a subspace of the environment. The nodes of a world graph are important waypoint states and edges represent feasible traversals between them. Our framework has two learning phases: 1) identifying world graph nodes and edges by training a binary recurrent variational autoencoder (VAE) on trajectory data and 2) a hierarchical RL framework that leverages structural and connectivity knowledge from the learned world graph to bias exploration towards task-relevant waypoints and regions. We thoroughly evaluate our approach on a suite of challenging maze tasks and show that using world graphs significantly accelerates RL, achieving higher reward and faster learning.
This paper proposes a novel approach to hierarchical reinforcement learning approach by first learning a graph decomposition of the state space through a recurrent VAE and then use the learned graph to efficiently explore the environment. The algorithm is separated into 2 stages where in the first stage random walk and goal conditioned policy is used to explore the environment and simultaneous use a recurrent binary VAE to compress the trajectory. The inference network is given the observation and action and the reconstruction is to, given the hidden state or hidden state+observation, reconstruct the action taken. The approximate posterior takes on the form of a hard Kumaraswamy distribution which can differentiably approximate a binary variable; when the approximate posterior is 0, the decoder must reconstruct the action using the hidden state alone. The nodes of the world graph are roughly states that are used to reconstruct the trajectories in the environment. After the graph is constructed, the agent can use a combination of high-level policy and classical planning to solve tasks with sparse reward.
Personally, I quite like the idea of decomposing the world into important states -- it is closely related to the concept of empowerment [1] which the authors might want to take a further look into. I believe extracting meaningful abstraction from the environment will be a key component for general purpose RL agent. One concept I really like in the paper is using the reconstruction error as the reward for the RL agent, which has some flavors of adversarial representation learning. Further, I also really like the idea of doing structured exploration in the world graph and I believe doing so can help efficiently solve difficult tasks.
However, I cannot recommend accepting this paper in its current draft as there might be potential major technical flaw and I also have worries about the generality of the algorithm. My main concerns are the following:
    1. The ELBO given in the paper is wrong -- the KL divergence should be negative. I want to give the paper the benefit of doubts since this could be just a typo and some (very rare) researchers use a reverse convention; however, this sign is wrong everywhere in the paper including the appendix yet the KL between Kuma and beta distributions uses the regular convention. I tried to check the source code provided by the authors but the code only contains architecture but not the training objective, training loops or environments. As such, I have to assume that the ELBO was wrongfully implemented, unless the author can provide the full source code, or, if the ELBO is indeed incorrectly implemented, rerun the experiments with the correct implementation.
    2. The proposed method for learning the graph does not have to be a VAE at all. The appendix shows that the paper uses a 0.01 coefficient on the KL, which is an extremely small value for VAE (in fact, most VAE’s have beta larger than 1 for disentanglement). Thus, I suspect the KL term is not actually doing anything; instead, the main reason why the model worked might be due to the sparsity constraints L_0 and L_T. In other words, the model is simply behaving like a sequence autoencoder with some sort of hard attention mechanism on the hidden code, which might explain why the model still worked well even with the wrong ELBO. To clarify, I think this is a perfectly acceptable approach for learning the graph and it would still be very novel, but the manuscript should be revised accordingly to reflect this. If the (fixed) VAE is important, then this comparison (0 KL regularization) would be a nice ablation regardless.
    3. Algorithm 1 requires navigating the agent to the key points from \mathcal{V}_p. This assumption is quite strong. When the transition dynamic is deterministic and fully reversible like the ones considered in the paper, using the reverse of replay buffer can indeed take the agent back to s_p, but in settings where the transitions are stochastic or the transitions are non-linear or non-reversible, how should the algorithm be used?
    4. It is not clear how \mathcal{V}_p are maintained. If multiple new nodes are added every iteration, wouldn't there be more than necessary nodes in \mathcal{V}_p? It seems to me some pruning criteria were used unless the model converged within small number of iterations? Are the older ones are discarded in favor of newer ones?
    5. How are the actions sequences “normalized”?
    6. In what way are the Door-Key environment stochastic? It seems like the other environments also have randomness, so is the only difference the lava pool?
I believe the propose method is sound, so if the revision can address either 1 or 2, I am willing to raise my score to weakly accept. If the revision in addition addresses 3, 4, 5, 6 in a reasonable manner, I am willing to raise my score to accept.
=======================================================================
Minor comments that did not affect my decision:
    - I think mentioning the names of the environment in the abstract might be uninformative since the readers do not know what they are a priori.
Reference:
[1] Empowerment -- An Introduction, Salge et al. 2014