Learning to Search from Demonstration Sequences

Dear Reviwer iWz4,

We express our gratitude for your thorough review and insightful comments on our paper. We are glad that you find our work interesting and the experimental results compelling. We address the concerns you raised as follows:

W1: Previous methods have introduced diverse differential network architectures to integrate different search algorithms into networks, as mentioned in the related works. It is unsurprising that integrating the best-first algorithm into the network has also yielded success.

We thank the reviewer for raising this question. There is a key distinction between the architectures mentioned in related works and our proposed method – they use a separately learned world model, while our method jointly learns a world model with the tree search. Moreover, our method constructs a search tree dynamically, instead of a fixed structure as in TreeQN, which allows it to search deeper compared to a shallow full tree, given the same computation budget.

In summary, we are the first work that proposes to differentiably construct a dynamic search tree and learn a world model jointly with the search algorithm. Further, we formulate the tree expansion as another decision-making task and solved it using REINFORCE. These unique properties set our work apart from existing related works.

W2: Discussion of Online RL

Thank you for raising this question. Our work primarily focuses on Offline RL settings and aims to learn to search from demonstration sequences only. The proposed method can be applied to an online setting without additional modifications. Although, in an online setting, the benefit of jointly learning a world model might be less significant, as the agent can explore unseen regions, collect data, and learn the policy. However, the proposed method would still benefit from its strong search inductive bias and improve the sample efficiency for training.

Q1: Do differences in the convergence rates of submodules exist? If present, could these differences impact the overall network's performance?

We do observe different convergence rates for different submodules. On the navigation problem, we observe that the reward function quickly converges after around 10k optimization steps, while the transition function and value function converge after around 40k steps. The difference in convergence rate may related to the difficulty of learning the submodule, for example, the reward function in the navigation problem is easy to learn as most rewards are 0 and only terminal states may have a positive reward. Generally, we observe that the performance of the overall network continues to increase even if some submodules have not converged yet.

We deeply value the thoughtful and constructive feedback you have provided. We sincerely hope we have addressed your concerns and raised your impression of our work. We are happy to clarify any other questions if required.

Dear Reviewer tL81,

Firstly, we express our gratitude for your thorough review and insightful comments on our paper. Your recognition of the novelty and contribution of our work is greatly appreciated. We have taken your feedback as an opportunity to further refine our manuscript. We address your concerns below:

W1: "Limited to deterministic Environments"

We agree that the current implementation of D-TSN is not aimed at stochastic environments and continuous action spaces. However, we briefly discuss how the current work might adapt to such scenarios below:

• For stochastic environments, we could take inspiration from existing works like [1], where we can incorporate an intermediate 'afterstate' in the search tree. In this approach, a state first deterministically transits to an intermediate 'afterstate', and then branches stochastically to the next state, which allows it to address the stochasticity in the environment.
• For continuous action spaces, search and planning methods often do not fit well into these problems. A popular approach is to discretize the action space so that it can fit into the search and planning framework. We could also perform multi-level discretization to trade-off between the depth and width of the search tree. For example, if a single action is defined by $k$ action variables $(a_1, .., a_k)$, we may discretize and branch on $a_1$ first, followed by discretization and branching on $a_2, ...,$ all the way to $a_k$, converting each action into a $k$-level action, and correspondingly increasing the depth of the tree. Our best-first search action selection will still result in a sparse tree for this scheme.

W2: "The computational complexity for the approach might be high".

We agree that constructing search trees is a computationally expensive operation, this is the nature of search-based approaches. However, search-based methods come with added performance improvements over model-free methods which is evident from our wide range of experiments. Notably, our method is more computationally efficient than other search-based methods such as TreeQN as we construct the tree by best-first search instead of full tree expansion. Moreover, in D-TSN, we can tune the performance-computation trade-off (by changing the number of expansion steps) depending on the computation requirements of the problem.

W3: "scalability is not thoroughly analyzed for longe-horizon tasks or higher dimensional state space"

Long-horizon tasks pose challenges such as delayed rewards, which may require doing more search iterations to look ahead further. This is a general challenge for reinforcement learning methods. Notably, the computational cost of our method scales linearly as the number of search iterations grows and offers a computationally-friendly approach to scaling to long-horizon problems compared to TreeQN (which scales exponentially in tree depth).

On the other hand, high dimensional states may require larger networks for each submodule in our framework. In terms of overall computational cost, the D-TSN framework grows by a constant factor as submodules become larger.

Q1. How to extend to stochastic decision-making problems?

Please see our response to W1.

Q2. How computationally expensive is D-TSN? How does it scale w.r.t. the action space size and the search tree depth?

We report the average time taken per training step and the total training time for a sample Procgen environment on a RTX 2080Ti GPU in the following table:

The computational cost of D-TSN grows linearly as the size of the action space increases. Since D-TSN uses best-first search, it does not have a fixed search tree depth. We quantify the complexity of the tree by the number of search iterations performed. Given this search budget, D-TSN may learn to search deeper or wider depending on the specific problem. Lastly, the computational cost of D-TSN grows linearly as the number of search iterations increases.

We sincerely hope that our clarifications above have increased your confidence in our work. We will be happy to clarify further if needed. We thank you again for sharing your valuable feedback on our work.

[1] Antonoglou, Ioannis et al. “Planning in Stochastic Environments with a Learned Model.” International Conference on Learning Representations (2022).

Dear Reviewer KW9a,

We thank you for spending time thoroughly reviewing our paper and providing constructive feedback, which helped us improve our work. We are pleased to know that you loved the motivation and overall idea behind our work. We address the concerns you raised as follows:

W1. Proof or empirical evidence on “slight changes in the network parameters could cause discontinuity in the tree structure”

We have provided the proof that 'the continuity of the Q-function is not guaranteed when the tree structure is not fixed' by giving a counter-example in Appendix B of the revised manuscript of the paper and have updated it in this submission. Please have a look. We are happy to provide further clarifications if needed.

W2: Fix grammatical errors and label every equation.

We thank you for pointing out these issues. We have performed a thorough grammar check and have labeled every equation in the revised manuscript.

W3.1: Why a search tree is a principal prerequisite for information retrieval?

The use of a search tree is adopted in various works, such as [1, 2, 3], in decision-making and planning problems. Search trees also show improvements in question answering-type information retrieval tasks, such as [4]. We are happy to discuss other 'IR methods' with the reviewer and will include them in the revised manuscript with specific references.

W3.2: I find it disingenious that the authors claimed that the search could be incomplete because the search process may visit regions unexplored during training. I think the reasoning here is incomplete and should be revisited by the authors.

Thank you for pointing out the issue of reasoning flow in paragraph 1-3. In paragraph 2, we stated that 'the separately learned world models are often incomplete' – by “incomplete” we mean that the world model is not trained on the complete state space and may perform badly on the states not seen during the training phase. Consequently, by using such world models, the search would suffer from the issue of compounding errors, particularly when the search is expanded into these unseen regions. In the next paragraph, we mentioned that end-to-end learning can improve the reliability of the world model. However, we might have not provided the clear reasoning in the previous manuscript. We provide the intended reasoning here: the advantage of end-to-end learning is that the world model and the search algorithm are trained together with a loss function that improves final performance of the resulting policy. In doing so, the search algorithm would learn to avoid going into state spaces where the world model performs badly (since going into these states would result in low final performance). We have added the above reasoning in the revised manuscript.

Q1: How is the full expansion mechanism in this work not a shallow tree ... since you mentioned that you have a fixed depth in your search protocol?

We believe there is a misunderstanding: we do not perform full expansion in D-TSN and we do not have a fixed depth. Instead, we perform best-first search to construct the tree for a fixed number of tree nodes (or search iterations), where the resulting tree structure can be deep and sparse as compared to a fully expanded shallow tree (as done in TreeQN).

Q2: When the encoder, value, reward, and transition functions are unavailable, how do they get jointly optimized with the search algorithms?

The encoder, value, reward, and transition functions are invoked by the search algorithm to finally derive the Q-values of the root node. Correspondingly, we use REINFORCE to train these submodules to minimize the expected loss on the Q-values. Please refer to Equation 16 in the revised manuscript for the loss function used by D-TSN. Also, please refer to Appendix D for a more visual-friendly illustration of the pipeline.

Q3: Why end-to-end learning improves the reliability of world models used in search?

Please see our response to W3.2.

Q4. Proof for "Dividing the network into these submodules reduces total learnable parameters and injects a strong search inductive bias into the network architecture, preventing overfitting to an arbitrary function that may align with the limited available training data."?

Firstly, in D-TSN, we construct a search tree by reusing submodules including encoder, value, transition, and reward function, this results in a smaller number of total network parameters as compared to constructing a tree of the same size without reusing the submodules.

In a discretized parameter space, we can apply Occam's Razor to derive a generalization bound proportional to the number of parameters. According to Occam's Razor principles, simpler models (those with fewer parameters) are favored, as they tend to generalize better to unseen data [1].

For continuous parameter space, we provide a proof sketch based on VC-dimension [5]. Consider a complex hypothesis space $H_{complex}$, a model from this space is flexible enough to fit any arbitrary functions, including those that align with the noise in training data. By dividing the network into submodules and reusing them multiple times in the search algorithm, we construct a constrained hypothesis space $H_{bias}$. The restriction of this bias leads to $VC(H_{bias}) < VC(H_{complex})$, which means $H_{bias}$ has a lower degree of freedom to fit arbitrary complex functions. Given $H_{complex}$ and $H_{bias}$, Occam’s Razor would favor $H_{bias}$ for its simplicity, as simpler models (those with lower VC dimension) tend to generalize better to unseen data. Thus, by injecting inductive search bias into the model design, we effectively reduce the risk of overfitting due to the decreased VC dimension and alignment with Occam's Razor principles.

Q5: Flow chart of the algorithm in the main text.

Thank you for this suggestion, we have added an illustration of the algorithm in the main text of the revised manuscript.

Q6: What metric is used in Table 1?

Quantities reported in Table 1 are success rate, i.e. whether the generated expression is valid, and results in the game of 24. We have updated the Table heading in the revised manuscript.

Q7: Appendix B, Lemma B.3: I'm sorry, is there supposed to be a running sum over all of x in the equation you wrote?

No, by Lemma B.1, B.2 and B.3, we show that h(x)=max(f(x), g(x)) is continuous, which is the building block of the Q-function represented by TreeQN – the backup phase calculates the value of each node by calculating the maximum Q-values of its children, which resembles h(x)=max(f(x), g(x)). We are happy to clarify any further questions.

Q8: Typos

Thank you for pointing it out. We have fixed these in the revised manuscript.

We sincerely appreciate the constructive critique you have provided and hope that our responses and the modifications to the manuscript adequately address your insightful feedback and increases your impression and confidence in our work. We are happy to provide any further clarifications if needed.

[1] Silver, David, et al. "Mastering the game of Go with deep neural networks and tree search." nature 529.7587 (2016): 484-489.

[2] Mohanan, M. G., and Ambuja Salgoankar. "A survey of robotic motion planning in dynamic environments." Robotics and Autonomous Systems 100 (2018): 171-185.

[3] Duarte, Fernando Fradique, et al. "A survey of planning and learning in games." Applied Sciences 10.13 (2020): 4529.

[4] Yao, Shunyu, et al. "Tree of thoughts: Deliberate problem solving with large language models." Advances in Neural Information Processing Systems 36 (2024).

[5] Blumer, Anselm, et al. "Learnability and the Vapnik-Chervonenkis dimension." Journal of the ACM (JACM) 36.4 (1989): 929-965.

Dear Reviewer KW9a,

We thank you for your thorough review of our paper and for providing constructive feedback that has significantly contributed to its improvement. Your insights have been invaluable in helping us refine our work.

We sincerely hope that our responses have sufficiently addressed the issues you highlighted in your review and follow-up comments. As the author-reviewer discussion period approaches its end, please do not hesitate to let us know if there is anything further we could do to improve your impression and final rating of our work.

Dear Reviewer 7dAQ,

We thank you for your insightful and constructive feedback on our manuscript. We appreciate your recognition of the novelty and practical application of our approach in addressing the complexity issue in TreeQN and its empirical evaluation. Below, we address your concerns to clarify and improve our work.

W1: How the approach handles large action spaces remains an empirical question

We agree with the reviewer that D-TSN can be expensive when the action space is large, as we need to evaluate every action’s Q-value. The computational cost grows linearly as the action space grows, this is a general challenge for value-based reinforcement learning methods when dealing with large action spaces. One potential way to adapt to a large action space is to have a Q-network that outputs the Q-values of all actions in one forward pass. Adapting D-TSN to use a Q-network to evaluate action branches is straightforward, and the search algorithm will remain largely the same.

W2: Application beyond discrete action space and deterministic environments uncertain.

We agree that the current implementation of D-TSN is not aimed at stochastic environments and continuous action spaces. However, we briefly discuss how the current work might adapt to such scenarios below:

• For stochastic environments, we could take inspiration from existing works like [1], where we can incorporate an intermediate 'afterstate' in the search tree. In this approach, a state first deterministically transits to an intermediate 'afterstate', and then branches stochastically to the next state, which allows it to address the stochasticity in the environment.
• For continuous action spaces, search and planning methods often do not fit well into these problems. A popular approach is to discretize the action space so that it can fit into the search and planning framework. We could also perform multi-level discretization to trade-off between the depth and width of the search tree. For example, if a single action is defined by $k$ action variables $(a_1, .., a_k)$, we may discretize and branch on $a_1$ first, followed by discretization and branching on $a_2, ...,$ all the way to $a_k$, converting each action into a $k$-level action, and correspondingly increasing the depth of the tree. Our best-first search action selection will still result in a sparse tree for this scheme.

Q1: Why choose REINFORCE, what about gumbel-softmax for differentiable sampling?

In the paper, we formulate the tree expansion as another decision-making task, and REINFORCE comes as a natural solution for such tasks. We agree that using the gumbel-softmax trick for differentiable sampling in Eq. 7 (in revised manuscript) can be a possible alternative. However, gumbel-softmax may require additional tuning of the temperature and dynamically adjusting it during training, which may make training more difficult. Nonetheless, it is an interesting approach that we can explore in the future.

Q2. Computation budgets of different methods?

We report the average time taken per training step and the total training time for a sample Procgen environment on a RTX 2080Ti GPU in the following table:

Q3. Writing suggestion.

Thank you for pointing it out. We have modified the sentence in the revised manuscript to “To construct the search tree, we employ a stochastic tree expansion policy and formulate it as another decision-making task. Then, we optimize the tree expansion policy via REINFORCE with an effective variance reduction technique for the gradient computation.”

Q4: If learning from demonstration sequences fails due to compounding errors due to the agent getting into states not during training, then the training distribution of the proposed method is also not sufficiently covered by the training distribution.

We agree with the reviewer that our method cannot cover more training distribution as compared to training the world model separately. The advantage of joint training is that the world model and the search algorithm are trained together with a loss function that optimizes the final performance. In doing so, the search algorithm would learn to avoid going into state spaces where the world model performs badly (since going into these states would result in low final performance). When the world model and the search algorithm are learned separately, the search would not avoid those unseen state spaces and thus may harm the performance.

Q5: What are the scores of the PPG sub-optimal policy?

The performance of the policy we used to collect trajectories is reported in the table below:

Q6. Std and error bars of Table 3?

We are in process of rerunning the experiments to calculate the std and error bars and will update once finished.

We sincerely appreciate the constructive critique you have provided and hope that our responses and the modifications to the manuscript adequately address your insightful feedback and increases your impression and confidence in our work.

[1] Antonoglou, Ioannis et al. “Planning in Stochastic Environments with a Learned Model.” International Conference on Learning Representations (2022).

Thank you for your prompt reply. We are pleased to hear that you find our response and revisions have added clarity to the manuscript.

Given the numbers in the table, would the authors think the performance difference between Model-free Q-network and D-TSN could be simply explained by the amount of computation?

This is a great question. We do not think the performance difference is merely due to computation. First, D-TSN is a search-based algorithm, which allows it to spend more inference-time computation by constructing a larger search tree. Model-free Q-network however, cannot scale at inference time as it simply gives the Q-values in one forward pass without constructing any structures that we can scale up. Model-free Q-network can be scaled up by increasing the model size and adding more training data, which requires more data collection and is harder to implement than inference-time scaling.

Compared with other search-based methods such as TreeQN, D-TSN spends an equal or less amount of computation while achieving better performance, which demonstrates the efficiency of D-TSN's search mechanism.