Kati8/ICLR_2023_2025 huggingface dataset:
authors: Dixant Mittal
authors: Liwei Kang
authors: Wee Sun Lee
reviews: 
Comments: This study presents a new differentiable tree search network comprising multiple submodules to estimate Q values. The network is optimized by minimizing the distance between the estimated Q values and the ground-truth Q values within a provided dataset. To address the substantial changes in the tree structure resulting from updates to the Q value function, a stochastic expansion policy is introduced to guide expansions in the search process. This policy ensures the continuity of the parameter space irrespective of changes in the tree structure.
Strengths: i) This work introduces the integration of the algorithmic inductive bias from the best-first search algorithm into a differentiable network, enabling automatic guidance of the tree search process through gradient backpropagation.

ii) This work underscores the significance of maintaining continuity of both the parameter space and the loss function which are dependent on the tree structure. To address this, the authors advocate for the adoption of a stochastic expansion policy to fulfill these prerequisites. 

iii) The experiment results are compelling. And it is particularly noteworthy to see the success achieved in tasks involving LLM.
Weaknesses: i) 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. Thus, it would be beneficial to compare the architectural variances between this method and previous methodologies.


ii) This work trains the overall network using an offline dataset. However, as extensively deliberated in preceding offline RL studies, this paradigm may get stuck when facing out-of-distribution states or actions. Thus, a comparative analysis between online training and offline training for the newly proposed network architecture could provide valuable insights.
Recommendation: 6
Reviewer confidence: 3
Soundness/correctness: 3
Clarity: 2
reviews: 
Comments: This paper presents a novel differentiable neural tree search architecture that can be learned directly from trajectories of data in a supervised manner. It essentially introduces a search-like inductive bias within the weights of the neural network. The proposed algorithm builds on TreeQN (Farquhar et al., 2018), and is crucially different since the proposed method allows building a tree structure that can stochastically sample from the action space, and not just be a fixed tree structure as in TreeQN. The authors empirically evaluate the strength of their approach by comparing against a Model-free Q network, a Model-based search method that trains the individual modules separately, and TreeQN, and report performance gains in various RL environments.
Strengths: - Builds on TreeQN and improves a significant limitation of the prior work, i.e., having a fixed tree structure. In reality, search algorithms should attempt to filter large action spaces and focus computation on promising variations in the tree. The proposed work gets around this limitation by sampling from the action space, and using REINFORCE to differentiate through the discontinuity of sampling.
- Strong empirical evidence that the proposed method improves on TreeQN, and having the modules trained separately.
- Strong ablation results showing the effectiveness of the proposed method.
Weaknesses: - Seeing how the approach handles large action spaces remains an empirical question, since currently, there is no policy that directly outputs a distribution over actions, instead the method requires the application of the transition network and the reward network for every action, which is not scalable to settings like, say, Go.
- The proposed method is mostly applicable to discrete action spaces with deterministic environments. Improving on this remains a future empirical question.
Recommendation: 8
Reviewer confidence: 4
Soundness/correctness: 4
Clarity: 2
reviews: 
Comments: This paper introduces the Differentiable Tree Search Network (D-TSN), a novel neural network architecture designed to learn search strategies from sequences of demonstrations without access to explicit search trees. D-TSN integrates the inductive bias of a best-first search algorithm into its structure, enabling the joint learning of essential planning submodules, including an encoder, value function, and world model. To construct the search tree, the authors employ a stochastic tree expansion policy, formulating it as a decision-making task optimized via the REINFORCE algorithm. They introduce a variance reduction technique using a telescoping sum to address high variance in gradient estimates. D-TSN is applicable in scenarios where the world model is known or needs to be jointly learned with a latent state space. The authors evaluate D-TSN on tasks from both scenarios, including the Game of 24, a 2D grid navigation task, and Procgen games. Experimental results demonstrate that D-TSN is effective, particularly when the world model with a latent state space is jointly learned, outperforming baselines in terms of success rate and generalization capabilities.
Strengths: 1. Novel Architecture: The paper proposes a novel neural network architecture, D-TSN, which embeds the inductive bias of a best-first search algorithm, allowing for the end-to-end learning of planning components from demonstration sequences. 

2. Joint Learning of Planning Components: D-TSN jointly learns the encoder, value function, and world model. This is advantageous when the world model is not given but is needed to be learned from data. 

3. Variance Reduction Technique: Authors use an effective variance reduction technique using a telescoping sum in the REINFORCE algorithm to addresses the high variance associated with policy gradient methods. 

4. Comprehensive Experiments: The method is applied to a wide variety of tasks, such as reasoning problems, navigation, and game environments, supporting the claim that it is versatile and effective across domains. 

5. Improved Performance: The authors show that D-TSN outperforms baselines, showing its problem solving performance in challenging tasks with limited supervision, especially in jointly learned world model settings.
Weaknesses: 1. Limited to Deterministic Environments: The current implementation is restricted to deterministic decision-making problems with discrete action spaces.

2. Computational Complexity: The computational complexity for the approach  might be high, because it consists of constructing search trees and performing REINFORCE updates. This can be a problem especially when applying for deeper trees or larger action spaces.

3. Scalability: scalability is not thoroughly analyzed for longe-horizon tasks or higher dimensional state space.
Recommendation: 10
Reviewer confidence: 4
Soundness/correctness: 4
Clarity: 4