Scaling Value Iteration Networks to 5000 Layers for Extreme Long-Term Planning

Q2:

In Figure 9, there is no plot of 100x100 maze. Why is such plot not provided? Given that the depth of 5000 is required only for such maze, it would be crucial to see how the choice of other depth values would affect the SR and OR.

Answer:

Thanks for this useful suggestion. We show the performance for $100 \times 100$ maze in Fig. 9 and Fig. 10 in the revised paper. For the convenience of the reviewer, we listed the performance in the table below. For tasks requiring more planning steps, it is crucial to use a network of depth 5000 to solve the task.

Table R3: The success rate for DT-VIN across various depths $N$.

Table R4: The optimality rate for DT-VIN across various depths $N$.

Q3:

The description to ViZDoom experiment is vague and more details are needed. Can the authors clarify how the optimal path can be found and what it means for “3D ViZDoom mazes with size M = 35”?

Answer:

In 3D ViZDoom, we follow the setting in the GPPN paper [r1], where the expert data are computed using Dijkstra’s algorithm on the underlying 2D maze map of the environment. For "3D ViZDoom with size $M = 35$ ", we refer to a maze designed on a grid of $35 \times 35$ cells. Each cell in this grid corresponds to an area of $64 \times 64$ map units within the 3D ViZDoom environment. The map unit is the basic measure of space used in the ViZDoom game engine to define distances and sizes. We have revised our paper to clarify this point.

[r1] Lee L et al. Gated path planning networks[C]//ICML, 2018.

Q4:

The evaluation on Rover Navigation is also unclear. The paper says that “elevation data are utilized solely for creating expert paths and assessing algorithm performance, not as input for the neural networks”. But the expert paths are needed to define the training loss (i.e., the additional loss), right? Then the proposed method DT-VIN still (implicitly) requires the evaluation data as the input?

Answer:

You're correct that in our experiments, elevation data is implicitly used in training. However, this is true for all the baselines, including our proposed DT-VIN and compared method, which use expert data during training, and, like them, DT-VIN does not use expert data during testing. Furthermore, even in the training phase, the elevation data is not strictly necessary for DT-VIN---only expert trajectories are necessary (a weaker requirement than methods needing the full dynamics inferred from elevation data, such as A*, would require). It's a caveat of this benchmark that we have access to elevation data and can use this to generate expert trajectories. Other methods like Monty Carlo exist to generate expert data without relying on such information and in line with that many benchmarks only provide expert trajectories. For example, the Mars Rover 2020 [r2] benchmark only provides expert trajectories built with programs written by expert rover drivers (we opted against this benchmark as it was not open source). We have clarified this point in the revised paper.

[r2] Ono, Masahiro, et al. "Data-driven surface traversability analysis for Mars 2020 landing site selection." 2016 IEEE Aerospace Conference. IEEE, 2016.

W3:

Further, the number of layers has a big impact on the loss scales. Consider the extreme case of 5000 layers, if the optimal path is 1000, then the final loss will have 4000+1 terms, which can make the optimization hard as the effective backpropagation of gradients can be challenging. The paper should discuss how this long-range backpropagation of gradients can be well handled in practice.

Answer:

Thanks for this very insightful comments. This is discussed in the appendix. Specifically, Appendix B.1, Line 728,

"To reduce the computational complexity of highway loss, we apply adaptive highway loss only to layers n satisfying the condition $n \mod J = 0$, where $J$ is a hyperparameter set to $10$ in our experiments."

Here, the main idea is to build the highway connections at interval $J$, for example, every 10 neural network layers. Using this, the number of the loss terms will reduce to only $1/J$ of the original one. As demonstrated in the table below, this mechanism with $J=10$ reduces training time to 1/3, while incurring a negligible performance degradation of just 0.01% compared to the configuration without it.

Table R2: Training Time and Success Rates (%) across Different Ranges of SPLs for DT-VIN with Different $J$ Values

Questions

Q1:

Since the training of highway connections requires the expert path from the start to the goal, how the training and testing mazes are constructed in the empirical evaluation? If the training and testing mazes are the same, the extra loss might leak the path info to the planning module.

Answer:

Thanks for this thoughtful comment. We would like to clarify that the highway loss is used only in the training phase, and not in the testing phase; after training, the learned model directly selects actions without again looking at the expert data.

In accordance with the experimental setup in the GPPN paper, the training and testing mazes are randomly generated using Depth-First Search with the Recursive Backtracker algorithm. We also enforce that all the maze maps, i.e., the positions of the walls and goals, in the dataset are different.

W2:

...highway loss relies on two important factors: the length of the optimal path (the provided expert path from the start to the goal) and the number of layers for the planning module. It would be important to discuss the situations where the optimal path might not be known in advance. For example, in the VizDoom experiment, it is unclear to me how such optimal path is obtained.

Answer:

Thanks for this insightful comment. Note that the knowledge of the length $l$ of the expert path naturally exists in the imitation learning case. However, for the case where such information is unknown, one can use either the length of non-expert data or some heuristic methods to estimate $l$ when the actual $l$ is completely unknown, e.g., using the distance between the start and the goal position. To further address your comment, we have now run a number of ablations on this.

First, to measure the effect of overestimation/underestimation, we conducted the experiment with various estimated values of the length of the shortest path $\widehat{l}$, which are $0, l/2, l, 2l, N$ (where $l$ is the actual length of the shortest path, $N$ is the depth of the planning module). Second, to evaluate the case when the estimation of $l$ has variance, we use $l \cdot \max(\epsilon, 0)$ as the estimation, with $\epsilon$ sampled from a Gaussian distribution $\mathcal{N}(1, 1)$. Third, we also assessed two additional variants for estimating $l$: (a) One variant utilizes the length of non-expert trajectories for $l$; (b) Another variant estimates the shortest path length heuristically using the L1 distance between the start $(x_s,y_s)$ and the goal $(x_g, y_g)$, i.e., $D=|x_s-x_g|+|y_s-y_g|$.

As indicated in the table below, both overestimation and underestimation led to a performance degradation of no more than 7%. Additionally, we found that leveraging non-expert data or the heuristic L1 distance only yields a nearly 3% degradation in performance, and performs better than the case when the optimal length is extremely overestimated/underestimated. These results imply that employing the information from non-expert data or heuristic estimation could be taken as an alternative when the optimal length is not available.

For the details in ViZDoom environment, please refer to our reply to Q3 below.

Table R1: The ablation study for using various estimated lengths of optimal paths for adaptive highway loss, under $35 \times 35$ ViZDoom navigation. The best results are highlighted by bold, while the second-best result is highlighted by $\underline{\text{underline}}$.

We appreciate the valuable comments and advice from Reviewer FzU7. Below, we provide detailed responses to the reviewer's questions.

Weaknesses

W1:

One of the main contributions claimed by this paper is the incorporation of observation into the transition kernel. However, as mentioned by the paper itself, “original VIN paper proposes a general framework where the latent transition kernel depends on the observation”. In other words, the VIN has incorporated the observation into the learning of transition kernel already, but its open-sourced implementation does not reflect this. So this can hardly be regarded as a contribution of this paper. The claim of incorporation of observation may sound like a good implementation suggestion when the considered the maze is of bigger sizes.

Answer:

There is some depth to the question of whether, proposing a method with a general definition but examining only a subset of that definition (e.g., invariant transition), allows another to claim a contribution by recognizing the importance of a different subset (e.g., observation-dependent transition). We would argue that the answer to this is yes. Indeed, we base our contribution on discovering that this specific version, i.e., observation-dependent transitions, allows for addressing scalability challenges not tackled by the original VIN paper. Moreover, the original VIN paper did not specify whether one should employ fully connected or convolutional layers for the transition mapping. In contrast, we advocate for the use of a single convolutional layer, which proves to be effective while adding minimal computational overhead. Altogether, we endeavour to be fair to the original work and thus acknowledge it here, but it does seem that we found something very important that was simply missed by the original VIN work: the specific variations that allow it to solve long-term planning. We believe this increases the practical applicability and robustness of the model, making it a substantial contribution beyond the original framework's scope. This adaptation opens new avenues for applying VIN in more diverse and challenging scenarios, which we have demonstrated with empirical results in our paper.

Q5:

Concerns about the novelty (observations-depdent, CNN).

Answer:

We think there has been a bit of confusion here and we apologize for not being more clear in our last response. Our proposed "dynamic transition kernel" goes well beyond the the "observation-dependent transition kernels" that are partly described in VIN. While observation-dependent transition kernels and the dynamic transition kernels we propose here both depend on the observation $x$ (formally, $\overline{\mathsf{T}}=f^{\overline{\mathsf{T}}}(x)$), the main innovation of our approach lies in the fact that our dynamic kernels are also a function of the latent state. In contrast, the observation-dependent kernel of VIN are not.

Formally, this means that (as shown in eq. (1) of the paper and reproduced below), the latent transition kernel ${\overline{\mathsf{T}}}_{i,j,\overline{a},i^{\prime},j^{\prime}}$ is different for each latent state $(i,j)$:

Dear reviewer FzU7, as the discussion period is nearing its end, please let us know if our recent rebuttal has successfully put to rest your concerns. We have made extensive efforts to address both your original and new comments, revising the paper as needed (in accordance with the expectations of the ICLR reviewing process). We understand you likely have many pressing responsibilities, and we appreciate the time you've already dedicated to reviewing our work. We are ready to respond if there are any remaining questions. Thank you once again for your valuable feedback.

Answer to Q3 (2/2):

c. Data: The scale of the dataset needs to scale up with the complexity of the problem rather than the model depth. Under the same scale of the problem, we didn't find that increasing model depth requires additional data. As shown in Table R9, without expanding the dataset, increasing the model depth does not reduce the performance.

Table R9: The success rate of DT-VIN across various dataset depths $N$, maintaining the same size as the original dataset.

Moreover, even in situations where data is rare, DT-VIN still outperforms compared methods. As shown in Table R10, with only 50% of the original dataset, DT-VIN greatly outperforms existing methods. We also highlight the changes compared to the performance with a full-sized dataset in Table R11, where DT-VIN results in less than a $0.2$ degradation for tasks within the range [1, 100], while the best-performing comparison method, GPPN, incurs a degradation of nearly $12$.

Table R10: The success rate for each method, using a dataset reduced to 50% of the original size.

Table R11: The changes in success rate for each method, using a dataset reduced to 50% of the original size, compared to the full-sized dataset (more negative is worse).

Q3:

... Concerns about scale:

a. Compute: ... How much GPU time was needed to train the 100x100 maze, 5000 layers?

b. Model size: What is the relationship between the number of planning steps required by the problem, and the depth of the network required to solve it? ... From the plots, it appears that success tails off as the number of steps approaches the network depth, but is it feasible to keep increasing network depth as the problems become more complicated?

c. Data: Is it necessary to keep scaling up the dataset in line with increasing the model depth, and does this limit the method’s applicability to other problem domains where training data can’t be so easily generated?

Answer to Q3 (1/2):

a. Compute: Table R7 shows the memory consumption and training time on an NVIDIA A100 GPU for DT-VIN and the compared methods when using 5000 layers and training for $90$ epochs on $100 \times 100$ maze. As shown in the table below, although DT-VIN takes more training time than GPPN, it requires much less GPU memory. These results are generally consistent with those observed in the $35 \times 35$ 2D maze in Table 5 of the original paper.

Table R7: Computational complexity of each method using 5000 layers and training for $90$ epochs, evaluated on $100 \times 100$ 2D maze navigation.

b.Model size: In our experiments, the depth of the network required to solve the problem is close to linear with the number of planning steps required by the problem. For maze size $M=15, 25, 35$, we test DT-VIN models at increasing depths in increments of 100 until the optimal performance is achieved. For instance, for mazes of size $25\times25$, we assess depths of $100,200,300,400$. For maze size $M=100$, we assess depths of $4000, 5000, 6000$. As Table R8 illustrates, the depth of the smallest network that can solve the task increases slightly more than linearly with the required planning steps. Therefore, it might be feasible to continue increasing the network depth as the problems become more complex.

Table R8: Minimal depths of DT-VIN model across various maze sizes.

W2:

I would have loved to see how DT-VIN performed in a very different domain. .....

Answer:

In response to your suggestion for diversity in our experimental domains, we have extended our evaluation to a continuous control setting, where both state and action spaces are continuous. Here, we use the Point Maze as the benchmark [r1], where the agent must decide the forces to apply to a controllable ball to reach a goal within 800 steps. We measure both the final distance to the goal and the success rate, where we say an agent has succeeded if the Euclidean distance to the goal is less than 0.5. As shown in the table below, DT-VIN achieved closer proximity to the goal and a higher success rate compared to the other methods.

Table R2: The performance of each method on continuous control tasks.

[r1] Fu, Justin, et al. "D4RL: Datasets for Deep Data-Driven Reinforcement Learning." arXiv preprint arXiv:2004.07219 (2020).

We sincerely appreciate Reviewer Pndq's constructive comments. Below we address the questions and concerns.

Weaknesses

W1:

What I feel the paper lacks is a motivation from a broader perspective: why do we need an RL agent that can solve bigger mazes? Particularly when – as acknowledged by the authors – the A* algorithm has been around since 1968 (and, indeed, Dijkstra’s algorithm is used in this work to compute the shortest path lengths).......

This is linked to my concerns about evaluation: all three domains (2D grid-world mazes, the 3D VizDoom, and the Rover navigation) feel like essentially the same problem. In particular, I can’t really see what value the VizDoom experiments add. ......

Answer:

Thank you for your comments regarding the motivation and evaluation in our paper.

First, we would like to note that traditional search-based algorithms such as A* and Dijkstra’s algorithm require a fully known and accurate environmental model to function effectively. This condition limits their applicability in scenarios where the MDP transitions or rewards are unknown or the state and action spaces are large or continuous. Our work instead considers environments without an accurate model, focusing on learning to generalize across unseen mazes directly from visual inputs rather than predefined graph structures. For example, in continuous control problems where the state and action space are continuous (as we will show in response to W2), classical search-based planning algorithms like A*, are ill-suited. In contrast, our model learns directly from observations to actions, demonstrating generalizability to unseen tasks.

While our tasks are fundamentally based on navigation and planning, they address distinct challenges that collectively showcase the broad applicability and robustness of our proposed method in both simulated and real-world scenarios.

The 3D ViZDoom environment (Sec 4.2) was selected to assess the planning abilities of our algorithms under conditions of perceptual uncertainty, where the input is derived from noisy, first-person visual inputs. This introduces a significantly different challenge that tests the algorithm’s effectiveness in environments similar to those encountered in robotics. The sentence "The maze designs for the 3D mazes are generated in exactly the same manner as the 2D mazes…" in GPPN just refers to the way the positions of obstacles are generated. As also mentioned by the reviewer, the input to the model of 3D ViZDoom is different from the 2D maze, and the model needs to predict a 2D maze, which could be noisy and bring difficulty to the planning module.

The Rover navigation task involves long-distance planning based on satellite imagery, presenting a practical challenge where short-term planning solutions are inadequate. This task is especially demanding due to complex real-world details such as variable lighting, textures, and the complications introduced by aerial photography.

Q2:

In the GPPN paper, three transition kernels were used (NEWS/Moore/Differential Drive) – were these considered for DT-VIN? ......

Answer:

Yes, we used the NEWS transition kernel for the maze environment because it's the standard one used in gridworld and maze environments, and our preliminary experiments didn't find a notable difference when swapping to different transition kernels. We mentioned this setting in Lines 269-270 in the original paper, the agent navigates the four adjacent cells by moving one step at a time in any of the four cardinal directions. We have run an additional ablation using different transition kernels. As shown in the tables below, DT-VIN consistently outperforms all the compared methods regardless of the kernel used.

Table R5: The success rate for each method in $35\times 35$ 2D maze navigation with Differential Drive transition kernel, where the agent can move forward along its orientation or rotate 90 degrees left or right.

Table R6: The success rate for each method in $35\times 35$ 2D maze navigation with MOORE transition kernel, where the agent can relocate to any of the eight adjacent cells that comprise its Moore neighborhood.

Questions

Q1:

The formulation of the highway loss requires knowledge of the shortest path $l$ (as acknowledged by the authors in Appendix A) – where this is unknown, what is the cost of over/underestimating it? ...

Answer:

Thanks for this insightful comment. Note that the knowledge of the length $l$ of the expert path naturally exists in the imitation learning case. However, for the case where such information is unknown, one can use either the length of non-expert data or some heuristic methods to estimate $l$ when the actual $l$ is completely unknown, e.g., using the distance between the start and the goal position. To further address your comment, we have now run a number of ablations on this.

First, to measure the effect of overestimation/underestimation, we conducted the experiment with various estimated values of the length of the shortest path $\widehat{l}$, which are $0, l/2, l, 2l, N$ (where $l$ is the actual length of the shortest path, $N$ is the depth of the planning module). Second, to evaluate the case when the etsimation of $l$ has variance, we use $l \cdot \max(\epsilon, 0)$ as the estimation, with $\epsilon$ sampled from a Gaussian distribution $\mathcal{N}(1, 1)$. Third, we also assessed two additional variants for estimating $l$: (a) One variant utilizes the length of non-expert trajectories for $l$; (b) Another variant estimates the shortest path length heuristically using the L1 distance between the start $(x_s,y_s)$ and the goal $(x_g, y_g)$, i.e., $D=|x_s-x_g|+|y_s-y_g|$.

As indicated in the table below, both overestimation and underestimation led to a performance degradation of no more than 7%. Additionally, we found that leveraging non-expert data or the heuristic L1 distance only yields a nearly 3% degradation in performance, and performs better than the case when the optimal length is extremely overestimated/underestimated. These results imply that employing the information from non-expert data or heuristic estimation could be taken as an alternative when the optimal length is not available.

Table RX: The ablation study for using various estimated lengths of optimal paths for adaptive highway loss, under $35 \times 35$ ViZDoom navigation. The best results are highlighted by bold, while the second-best result is highlighted by $\underline{\text{underline}}$.

The extra work over the last week or two has, I believe, clearly demonstrated the proposed method's superiority over the previous baselines when it comes to other environments (eg continuous action space), data efficiency, and different transition kernels. The authors have also successfully shown that the method's reliance on knowing in advance the length of the expert path (l) isn't the weakness it first appeared to be. In the light of this, I'm revising my assessment of the paper's contribution, and my overall rating.

W1.3:

Although adaptive highway loss was introduced to address gradient stability, the ablation study suggests that the softmax latent transition kernel plays a larger role in mitigating gradient issues. Some theoretical justifications or more explanation of this gap would be helpful. Including loss curves/gradients norm graphs on one or more of the environments from the paper would help visually illustrate the trend of gradient instability for lack of each ingredient. ...

Answer:

Thanks for this insightful comments. The softmax operation ensures that the values of the dynamic transition kernels remain within $[0,1]$, helping to prevent the gradient exploding problem. In our experiments, we found that the gradient of DT-VIN lacking softmax operation explodes at the first forward-backward pass of training, resulting in the loss escalating to NaN (Not a Number) during the training process. Tables R2 and R3 shows the gradient and the loss of with and without Softmax Operation.

The adaptive highway loss improves gradient flow toward shallower layers. As shown in Tables R4 and R5, without adaptive highway loss, the L1 norm of the gradient for DT-VIN is closer to zero in the first 10 layers of the network. The adaptive highway loss can reduce this vanishing gradient problem, resulting in lower loss.

We have also presented the figures of these results in Section C.4, Figures 12 and 13 in the revised paper.

Table R2: The L1 norm of gradient averaged over the first 10 layers during the training process for DT-VIN with and without Softmax Operation.

Table R3: The loss during the training process for DT-VIN with and without Softmax Operation, evaluated on $35 \times 35$ 2D maze with depth $N=600$.

Table R4: The L1 norm of gradient averaged over the first 10 layers during the training process for DT-VIN with and without Adaptive Highway Loss.

Table R5: The loss during the training process for DT-VIN with and without Adaptive Highway Loss.

W2:

Testing on additional environments, such as continuous control domains or graph-based planning tasks, would provide a more comprehensive evaluation of DT-VIN. The original VIN paper included both types - learning policies with guided policy search with unknown dynamics and simulating in Mujoco for the continuous control domain (metric is measuring the final distance to the target....

Answer:

In response to your suggestion for diversity in our experimental domains, we have extended our evaluation to a continuous control setting, where both state and action spaces are continuous. Here, we use the Point Maze as the benchmark [r1], where the agent must decide the forces to apply to a controllable ball to reach a goal within 800 steps. We measure both the final distance to the goal and the success rate, where we say an agent has succeeded if the Euclidean distance to the goal is less than 0.5. As shown in the table below, DT-VIN achieved closer proximity to the goal and a higher success rate compared to the other methods.

Table R6: The performance of each method on continuous control tasks.

[r1] Fu, Justin, et al. ”D4RL: Datasets for Deep Data-Driven Reinforcement Learning.” arXiv preprint arXiv:2004.07219 (2020)

W1.2:

The adaptive highway network depends on an estimation of ($l$) (the trajectory/planning path length). It would be insightful to include robustness analysis of the method against different ($l$) values, exploring how sensitive DT-VIN’s performance is to the estimation of ($l$). For example, how does the success rate change in Vizdoom when you overestimate/underestimate ($l$), or when the estimate of ($l$) has high variance? The metric is the success rate.

Answer:

Thanks for this insightful comment. Note that the knowledge of the length $l$ of the expert path naturally exists in the imitation learning case. However, for the case where such information is unknown, one can use either the length of non-expert data or some heuristic methods to estimate $l$ when the actual $l$ is completely unknown, e.g., using the distance between the start and the goal position. To further address your comment, we have now run a number of ablations on this.

First, to measure the effect of overestimation/underestimation, we conducted the experiment with various estimated values of the length of the shortest path $\widehat{l}$, which are $0, l/2, l, 2l, N$ (where $l$ is the actual length of the shortest path, $N$ is the depth of the planning module). Second, to evaluate the case when the etsimation of $l$ has variance, we use $l \cdot \max(\epsilon, 0)$ as the estimation, with $\epsilon$ sampled from a Gaussian distribution $\mathcal{N}(1, 1)$. Third, we also assessed two additional variants for estimating $l$: (a) One variant utilizes the length of non-expert trajectories for $l$; (b) Another variant estimates the shortest path length heuristically using the L1 distance between the start $(x_s,y_s)$ and the goal $(x_g, y_g)$, i.e., $D=|x_s-x_g|+|y_s-y_g|$.

As indicated in the table below, both overestimation and underestimation led to a performance degradation of no more than 7%. Additionally, we found that leveraging non-expert data or the heuristic L1 distance only yields a nearly 3% degradation in performance, and performs better than the case when the optimal length is extremely overestimated/underestimated. These results imply that employing the information from non-expert data or heuristic estimation could be taken as an alternative when the optimal length is not available.

Table Rx: The ablation study for using various estimated lengths of optimal paths for adaptive highway loss, under $35 \times 35$ ViZDoom navigation. The best results are highlighted by bold, while the second-best result is highlighted by $\underline{\text{underline}}$.

We sincerely thank reviewer pQwz for the valuable comments. Below we address the questions and concerns.

W1.1:

Additional ablation studies would clarify the benefits of directly connecting hidden layers to the final loss instead of intermediate layers. Furthermore, the claim that having ( $n > l$ ) is advantageous for short-term planning tasks needs stronger empirical backing. It would be beneficial to include more controlled experiments to justify this claim. It would strengthen the claim if there is a performance gap between the two methods in 35x35 and 100x100 2D maze environments in the “shorter SPL” regime. i.e., if you do not have ( $n > l$ ) in your loss, does the performance drop in the [1,100] regime for 35x35 and [1,600] regime for 100x100? To be clear, I am asking for the 2D maze environments to be tested under these different regimes with success rate as the metric as in the paper.

Answer:

Following your suggestion, we conducted an ablation study for adaptive highway loss by evaluating the following variants in shorter planning tasks:

1. Implement skip connections for intermediate layers of the planning module of VIN, like what has been done in Residual Nets [r1] and Highway Nets [r2]. As shown in the table below, this variant performs poorly, achieving only 61.35% success rate in comparison to DT-VIN's 99.98% on $100 \times 100$ Maze. These results are consistent with those in existing work [r3].

2. $\mathbf{1}_{ \lbrace n = l \rbrace }$, only building highway loss for a specific layer $n$ which satisfies $n=l$. As shown in the table below, this variant performs worse than the adaptive highway loss, showing that the component $n>l$ plays an important role in the performance.

3. Building highway loss for all intermediate layers $n$, without the term $\mathbf{1}_{ \lbrace n \ge l \rbrace }$. This variant is already verified to be less effective in the original paper, which is called "Fully Highway Loss" (see Figure 4 (c) or Section 4.1, Ablation Study, Adaptive Highway Loss).

Table R1: The success rates for the variants of adaptive highway loss.

[r1] He K, et al. Deep residual learning for image recognition[C]. CVPR. 2016.

[r2] Srivastava R K, et al. Training very deep networks[C]. NeurIPS, 2015.

[r3] Wang Y, et al. Highway Value Iteration Networks[C]. ICML, 2024.

Dear reviewer pQwz, as the discussion period is nearing its end, please let us know if our recent rebuttal has successfully put to rest your concerns. We have made extensive efforts to address both your original and new comments, revising the paper as needed (in accordance with the expectations of the ICLR reviewing process). We understand you likely have many pressing responsibilities, and we appreciate the time you've already dedicated to reviewing our work. We are ready to respond if there are any remaining questions. Thank you once again for your valuable feedback.

We appreciate the insightful and helpful comments from Reviewer vguh. Please find our responses to your concerns and questions below.

Weaknesses

W1:

Clarification of Motivation: While the motivations for introducing a dynamic transition kernel and an adaptive highway loss are clearly explained, the significance of investigating the Value Iteration Network (VIN) remains unclear. In the introduction, the paper notes that "Within reinforcement learning, one notable method is the Value Iteration Network (VIN)," but the paper does not expand on why VIN is particularly notable or important to study, especially in comparison to other established methods like Dreamer and MuZero. Adding a discussion that situates VIN’s relevance among these methods could strengthen the motivation and context for readers.

Answer:

Thank you for the constructive feedback. We have revised the introduction in Section 1 to better articulate the unique position and advantages of VIN within the spectrum of AI planning techniques. VIN stands out in the reinforcement learning landscape due to its unique architecture, which integrates the latent planning module into the deep neural network. Unlike Dreamer, which leverages learned models to simulate or predict future states [r1], and MuZero, optimizes policies through a sophisticated tree-based search strategy using predictions from its learned model [r2]. VIN's integrated architecture allows it to be trained end-to-end, meaning that both the state representation components and the planning components are trained simultaneously. Moreover, the integrated planning capability of VINs enables them to effectively generalize to related but unseen tasks.

[r1] Hafner, Danijar, et al. "Dream to control: Learning behaviors by latent imagination." ICLR, 2020.

[r2] Schrittwieser, Julian, et al. "Mastering atari, go, chess and shogi by planning with a learned model." Nature, 2020.

W2:

Clarity of Equation 2: It would be helpful to present the overall loss function in addition to the adaptive highway loss, as showing only $\mathcal{L}$ here may introduce ambiguity. Clarifying the meaning of x and y in Equation 2 is also necessary.

Answer:

We thank the reviewer for the suggestion and have updated the paper accordingly. Eq. 2 aims to provide a general form of $\mathcal{L}$ with arbitrary single-term loss function $\ell$, which depends on the task. For example, in imitation learning with discrete action space, where $x$ is the observation, the label $y$ is the expert action, and $\ell$ is the cross-entropy loss, the overall loss function can be rewritten as

where $f^{\pi} \left( \overline{V}^{(n)}(x), a \right)$ represents the probability of taking action $a$ given observation $x$.

For reinforcement learning where the policies are learned from reward feedback, the overall loss function can be rewritten as