Learning to Search from Demonstration Sequences
Authors: Dixant Mittal
Liwei Kang
Wee Sun Lee
Venue: ICLR 2025 Oral (2025)
Keywords: planning
reasoning
learning to search
reinforcement learning
large language model
Primary area: reinforcement learning
TL;DR: We propose a method that constructs search tree in a differetiable manner, and can be trained from just demonstration sequences.

Abstract: Search and planning are essential for solving many real-world problems. However, in numerous learning scenarios, only action-observation sequences, such as demonstrations or instruction sequences, are available for learning. Relying solely on supervised learning with these sequences can lead to sub-optimal performance due to the vast, unseen search space encountered during training. In this paper, we introduce Differentiable Tree Search Network (D-TSN), a novel neural network architecture that learns to construct search trees from just sequences of demonstrations by performing gradient descent on a best-first search tree construction algorithm. D-TSN enables the joint learning of submodules, including an encoder, value function, and world model, which are essential for planning. 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. D-TSN can be applied to problems with a known world model or to scenarios where it needs to jointly learn a world model with a latent state space. We study problems from these two scenarios, including Game of 24, 2D grid navigation, and Procgen games, to understand when D-TSN is more helpful. Through our experiments, we show that D-TSN is effective, especially when the world model with a latent state space is jointly learned. The code is available at https://github.com/dixantmittal/differentiable-tree-search-network.

Published as a conference paper at ICLR 2025
LEARNING TO SEARCH FROM DEMONSTRATION SEQUENCES
Dixant Mittal
Liwei Kang'
Wee Sun Lee
National University of Singapore
Moovita
[ixxnnt, kang, leews] @comp nus edu. sg
ABSTRACT
Search and planning are essential for solving many real-world problems. How- ever, in numerous learning scenarios, only oction-observation sequences, such as demonstrations or instruction sequences, are available for learning. Relying solely on supervised learning with these sequences can lead to sub-optimal performance due to the vast, unseen search space encountered during training. In this paper, we introduce Differentiable Tree Search Network (D-TSN), novel neural network architecture that learns to construct search trees from just sequences of demon- strations by performing gradient descent on a best-first search tree construction algorithm. D-TSN enables the joint learning of submodules, including an encoder, value function, and world model, which are essential for planning. To construct the search tree, 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. D-TSN can be applied to problems with a known world model or to scenarios where it needs jointly learn a world model with latent state space. We study problems from these two scenarios, including Game of 24, 2D grid navigation, and Procgen games, to understand when D-TSN is more helpful. Through our ex- periments, we show that D-TSN is effective, especially when the world model with a latent state space is jointly learned. The code available at httttttt////tthh/ comm http://aaannnttaaafefeeeeeeeeeee
INTRODUCTION
Search and planning are critical components in a wide range of complex tasks. It has long been important in areas such as robotics and games (Mohanan & Salgoankar 2018 Duarte et al. 2020).
and has recently gained prominence in the area large language models as inference-time solutions
for improving reasoning and agentic abilities (Yao et al. 2023 Hao et al. 2023 Zhou et al. 2024).
In many of these domains, step-by-step demonstrations are available in the form of recorded expert demonstrations, medical treatment records, and instruction/solution manuals. However, constructing a search tree would require data that is not directly available from the demonstration sequences this can be impractical or very expensive obtain. Hence, would be desirable learn to search from only demonstration sequences.
Search and planning can be effectively executed if simulator or world model available. When such a world model is not available, we may learn the world model and then use it for these pur- poses (Kaiser al. 2020 Hainer et al., 2020 Ha Schmidhuber, 2018). However, when separately learned from demonstration sequences, such world models often suffer from compounding errors, and cannot cover the complete state space. Using such world models for search can be unreliable. particularly if the search expands into parts of the space that are not seen during training (Asadi et al. 2018).
One way to increase the reliability of the world model is to perform end-to-end learning (Farquhar et al. 2018). where the world model trained with the search algorithm. One of the advantages of end-to-cee end learning that the search algorithm will learn to avoid going into states where the world model is inaccurate since this would result in lower performance. We begin by showing that, when the search tree has a fixed structure, the loss function for training is continuous with respect to the

---

Published as a conference paper at ICLR 2025
model parameters, justifying the use of fixed-structure search tree with gradient-based learning methods. However, if the search tree structure can change during training, the loss function may no longer be continuous. This is because small change in parameters could result in the construction of a different search tree and bring a large change in the corresponding loss function. To alleviate this issue, we propose employing stochastic tree expansion policy and optimizing the expected loss, ensuring the continuity of loss function in the parameter space.
We further propose formulating the search tree expansion as another decision-making task with the goal of minimizing the prediction error progressively during the tree expansion and refining the tree expansion policy via REINFORCE (Williams) 1992). The use of the REINFORCE algorithm introduces another challenge; REINFORCE usually has high variance in its gradient estimates. To handle that, we propose the use of baseline for variance reduction, adopting the telescoping sum trick used in Guez et al. (2018) for constructing the baseline.
Employing these techniques allows us to train novel neural network, Differentiable Tree Search Network (D-TSN), that embeds the structure of best-first search algorithm into its architecture. D-TSN can use a simulator or world model, if available, to learn to perform search. We study this scenario by fine-tuning a large language model to play the Game of 24 Yaoo et al. (023). D-TSN can also be used in more difficult scenarios where the world model is not available and only sequences of demonstrations navigation problem and Procgen games (Cobbe et al. 2020) which are substantially more complex. Through our experiments, we show that end-to-end learning for doing search and planning from only sequences of demonstrations is possible and useful, especially when the world model needs to be jointly learned.
provided. We study two problems in this scenario, small-scale 2D grid
RELATED WORKS
Search and planning are crucial in various domains. For example, in language model reasoning, recent works integrated search algorithms to enhance the reasoning capabilities of language models. These techniques employ in-context queries Yao al. [023]] Xie et al. [2033] Hao et al. |023] Zhou et al. 2024) or additional value heads (Liu et al. 2023) to guide the search process during inference, enhancing model's reasoning ability by looking at different possible solutions. In gaming, notable successes (Silver et al., 2016; Schrittwieser et al. 2020; Nasiriany et al. 2019) have combined deep
neural networks with search algorithms to achieve superhuman performance.
One line of recent works has focused on embedding search inductive bias into network architectures. Neural Admissible Relaxation (NEAR) (Shah et al. 2020) develops approximately admissible heuris- ties for the A algorithm. Neural A Search (Yonetani et al. 2021) integrates the A algorithm into network structures, learning cost function from gridworld map. Similarly, MCTSnets (Guez et al., 2018) incorporate the framework of Monte Carlo Tree Search (MCTS) into network architectures, guiding the search using parameterized memory embeddings stored in a tree structure. However, these methods rely on a known world model for planning. which limits their application when the world model is unknown. Another notable work, TreeQN (Farquhar et al. 2018). incorporates search inductive bias into the network by fully expanding search tree to a fixed depth while jointly learning a world model, which allows it to tackle problems where the world model is unknown. However, the full expansion mechanism of TreeQN leads to a shallow search tree, which limits its application in complex problems with longer trajectories.
Our setting of learning from sequences of demonstrations is closely related to Offline Reinforcement Learning (Prudencio et al., 2024, Levine et al., 2020), where an agent learns its policy solely from fixed dataset of experiences without further interactions with the environment. The Offline- RL paradigm is appealing in many real-world applications such as education, healthcare, and robotics, where active data collection is often infeasible (Singh et al. 2022 Singla et al., 2021 Liu al.) 2019). In our work, we develop a method that learns to search in an end-to-end fashion solely from sequences of demonstrations, such as text solutions to reasoning problems and expert trajectories on navigation and games.
3 DIFFERENTIABLE TREE SEARCH NETWORK
Differentiable Tree Search Network (D-TSN) is neural network design that incorporates the algorithmic inductive bias of a best-first search algorithm into the network structure. It learns from sequences of demonstrations to construct search trees by composing submodules, that include an
2

---

Published as a conference paper at ICLR 2025
encoder, value function, reward function, and transition function. If some submodules readily available, c.g. a world model or simulator, they can be directly used in D-TSN. When these submodules are not available, they can be jointly optimized with the search algorithm. This joint optimization allows the learned imperfect world model to be useful for online search and makes the submodules robust against errors in the world model. We start by describing the version of D-TSN where all submodules are learned.
3.1 LEARNABLE SUBMODULES
D-TSN comprises several learnable submodules that function as subroutines in a best-first search algorithm and dynamically construct computation graph. An illustration is provided in Figure The Encoder module (Eg) transforms the actual state 8, into latent state h, = (xt)), facilitating online search within a latent space. The Transition module (Tp) approximates the environment's transition function, using he and action ae to predict the subsequent latent state, h4-1 = To (he az), of the transition. The Reward module (R₂) approximates the environment's 's reward function. predicting the reward, r3 = Ro( az), for the transition based on h2 and a, The Value module (V₂) maps latent state h its estimated state value, Vo (h+). 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.
3.2 TREE SEARCH IN LATENT SPACE
The search begins by encoding the input state so into its latent state hgg then proceeds through expansion and backup phases. During expansion, the search tree expands iteratively for a set number of expansion steps, where each step expands node in the search tree. During backup. Q-values at the root node are recursively computed using the Bellman equation across expanded nodes. Each node represents a latent state reachable from the root, and branches represent actions taken. A candidate set O is maintained during the expansion phase, representing nodes eligible for further expansion. The D-TSN algorithm is detailed in Algorithm
Expansion Phase Each search iteration begins by evaluating the path value, V (N), of the candidate nodes. The path value defined as the cumulative sum of rewards from the root node to particular leaf node N, in addition to the value of the leaf node predicted by the value module (V)), i.e.
(NN)= = Rho oo) + + (oor
A naive implementation of the search selects the node N with the highest total path value for expansion; however, for differentiable search, the node Ñ is sampled from the candidates using a distribution constructed by applying softmax over the path values of the candidates. Expansion of node N is carried out by simulating every action on the node using the transition module (Tp). Simultaneously, the associated reward, Ro xr a)) is also computed. The resulting latent states are added to the tree as children of node N. Additionally, they are added to the candidate set O for subsequent expansions, while Ñ is excluded from the set. This can be represented as:
O - OU (hal ha To a): Va E A} Ñ
Backup Phase The expansion phase is followed by the Backup phase. In this phase, values of all tree nodes are recursively updated using the Bellman equation as follows:
Q(N, a) Ro(hN,,)) + V(N'),
where N' a))
(3)
V(N")- = [Volbxr), max, Q(N',a).
N' is a leaf node otherwise
(4)
After the backup phase, Q-values at the root node are returned as the final output of the online search. An illustration of the constructing the tree is shown in Figure]
3.3 CONSTRUCTION OF COMPUTATION GRAPH
Throughout the expansion and backup phases, illustrated in Appendix Figures 5 and [ dynamic computation graph is constructed where the output Q-values depend on the combination of all the submodules, i.e. Encoder, Transition, Reward, and Value modules. During training, the output
3

---

Published as a conference paper at ICLR 2025
RL
O2
Expansion
Backup
(1
V22
NOFFOO
Volue
O F (O..)
?IO (O. +1
V/ max. RIO.
Expansion
Backup
Sample leaf
nodes updated Q value of children nodes
softmax(y (O). to expand
Figure 1: An illustration of the stochastic tree expansion phase and backup phase.
Q-function evaluated by a loss function and optimized using gradient-ban based optimizers such as Stochastic Gradient Descent (SGD). These optimizers backpropagate the gradient of this loss through the entire computation graph and update parameters submodules an end-to-end process. An illustration of the process is provided in Appendix Figure
3.4 DISCONTINUITY OF THE Loss FUNCTION
pivotal aspect of optimizing D-TSN's parameters through gradient descent is ensuring that the loss function applied to the output Q-values is continuous in the network's 's parameter space. We initiate this discussion by showing that the loss function, when applied to the Q-values computed by expanding search tree, continuous in the parameter space, provided that the search tree has fixed structure.
Theorem 3.1. Given set of parameterized modules that are continuous the parameter space the Q-values computed by expanding search tree with fixed structure, through the composition of these modules and backpropagation of child values using addition and max operations, are continuous in the parameter space 0. When the tree structure is not fixed, the continuity of the Q-values is not guaranteed. (For a detailed proof, refer to Appendix B
Theorem 3.1 establishes that the Q-values calculated by fixed-structure tree expansion continuous in the network's s parameter space. The theorem applies to TreeQN (Farquhar et al. 2018), which performs full tree expansion output Q-values also continuous, contributing to TreeQN's success gradient-based optimization.
fixed depth, and consequently, the loss function applied to its
However, in the case of naive implementation of D-TSN, outlined in Section 3.2 the search tree is constructed by expanding only those paths that are likely to represent the optimal trajectory from the root node. This results in the output Q-values and the corresponding loss function being dependent on the specific tree, constructed during the search:
L(s,a) ((Qu(s,))))
(5)
When the network parameters are changed slightly, this naive implementation could generate different tree structure, which, in turn, would induce large change in the loss function and affect its continuity the parameter space.
3.5 ENABLING CONTINUITY IN THE LOss FUNCTION
To overcome the discontinuity issue observed in the naive implementation of D-TSN, we employ a stochastic tree expansion policy. This approach allows us to optimize the expectation loss function, defined as:
L(s, a) -E -E.
Σ ((())c((o(s..))))
The expected loss in Equation 6 is continuous in the parameter space 0 and can optimized using gradient-based optimization techniques.
3.6 STOCHASTIC TREE EXPANSION POLICY
In order to compute the expected loss in Equation (6). let us represent a partial search tree after iterations 7:- The output Q-vuuuss denoted Q0(s, al7). depend the final tree sampled after
4

---

Published as a conference paper at ICLR 2025
T iterations. We define stochastic tree expansion policy TT( that takes tree 7i input and outputs a distribution over the candidate nodes, facilitating stochastic selection of the node for further expansion and generating the tree 7++1- We compute the stochastic tree expansion policy by taking softmax over the path value, as defined Equation ① of each candidate node as follows:
exp(V(n))
TT(n)))) E (xp')(n"))
The gradient of the expected loss (Schulman et al., 2015) in Equation (6 can be computed as follows (See Appendix [foo the derivation):
Val(s, E. ((((((0().
v. + Vol((o((,a
(8)
3.7 REDUCING VARIANCE USING TELESCOPIC SUM
The REINFORCE term of the gradient in Equation 8 usually has high variance due the difficulty of credit assignment (Pignatelli et al. 2024 in reinforcement learning type objective; the second part of the gradient equation the usual optimization of loss function, so we expect it be reasonably well behaved. To reduce the variance of the first part of the gradient, we take inspiration from the telescoping sum trick used Guez al. 2018).
Let us denote the loss after l. iterations as L, (s, a.T:)). The objective is to minimize the loss after iterations, represented as LT. Assuming that Lo O, we can rewrite LT as a telescoping sum
LT = LT Lo L
Now, we define reward term, Yt, for selecting node during the 1 iteration as the reduction in the loss value after the thh iteration, i.e. Tu TT-1- Further, let us represent the return (or rewards-to-go from iteration to the final iteration T as RG, which can be computed as:
R₁ Σ Ti LT Lr-1
(10)
Given this, the REINFORCE term from Equation 8 can be reformulated as Σ Va log To (nglm)) Rs where LT serves as baseline that helps in reducing variance. Consequently, the final gradient estimate of the loss in Equation (6) is expressed as:
VoL(s,a) = E, Σ (o) VoL(Qo(s, (
(11)
For empirical evaluations, we use a single sample estimate Schulman et al. (0155) of the expected gradient in Equation (II).
3.8 Loss FUNCTIONS
In this section, we define series loss functions that are used to train D-TSN. Given a dataset of
trajectories where each trajectory is denoted as {(s,, o; ri, Q.) the primary objective is to make the computed Q-values closely approximate the observed Q-values for corresponding states and actions. To achieve this, we minimize the mean squared error between the predicted and observed Q-values. This loss, denoted as Lo, is defined as:
Le [s ((s.a.))-(Qa(s.. a) Q)²
(12)
To avoid Q-values for diitr---istiiittonn actions getting overestimated, incorporate the CQL (Ku
mar et al. 2020) loss, which encourages the agent to adhere to actions observed within the training data distribution. This loss, Lo. defined
LD = E(s.a)DD log Exp((Qoss, d')) - Qo(s, a))
(13)
5

---

Published as a conference paper at ICLR 2025
Additionally, we incorporate self-supervised consisteney loss functions (Schwarzer et al. 2021
Ye
et al., (022)) to ensure consistency in the transition and reward networks. Consider actual states s; and 11+1. where 84+1 is obtained by taking action a, in state 8;- Their corresponding latent state representations are denoted as h2 and h1+1- Here, h2 = Eo and hi+1 = (8(+1)) Now, we can use the transition module to predict another latent representation of state 81+11, represented as hu+1 = To (h;, 0z). To ensure that the transition function, Te. provides consistent predictions for the transitions in the latent space, we minimize the squared error between the latent representations hi+11 and ho41-
CT B(s, ons 1)) h++1)²
(14)
In a similar vein, we minimize the mean squared error between the predicted reward Ro a.) and the actual reward observed ri in the training dataset D.
CR, Bee,, de, ()))D((RA)) (Ro(hs, ai) - rz)²
(15)
Combining Equations
12
and 15 we define the loss function for training D-TSN as
DDTSS = E AxLC + AeCD] + ALCT. + XLCRa
(16)
where denotes the search tree. Additional details these loss functions are presented in Ap-
pendix D.4
4
EXPERIMENTS
4.1
D-TSN WITH KNOWN WORLD MODEL
We take an initial step towards learning to perform tree search from instruction sequences for LLMs by first demonstrating that LLMs can be effectively fine-tuned when the world model is known. We study Game of 24 (Yao al. 2023), a reasoning problem where the transition dynamics and rules are well-defined and readily avai able. The task is that given four cards drawn from deck of poker, use basic arithmetic operations to combine the card values to reach 24. The problem can be formulated as follows: a state is list of current numbers, initially 4; the actions involve selecting an operation as well as two numbers from the current numbers; the transition will calculate the operation and update the list of current numbers by removing the selected two and adding the calculated result; the reward is one when 24 is the only number in the list, and zero otherwise. Since language models can directly process token sequences,
do not need an Encoder in this problem. We use large language model, Llama3-8B (Dubey et al. 2024). with an additional value head as the value function Vo-
4.1.1 BASELINE
We compare this version of D-TSN with a supervised fine-tuned (SFT) language model. For fair comparison, the supervised fine-tuned model is also allowed to perform search during inference, where we use the log probability of the model generating the sequence as its value:
N-1
1 N
(s()))
log p(s41111))1...))
(17)
where s; are tokens in the sequence, p(-1-) is the conditional probability of generating the next token. The log probability is divided by N to normalize the effect of sequence length.
4.1.2 TRAINING DATASET
We collected all valid Game of 24 problems and their solutions through an exhaustive search of all combinations, then randomly selected 530 problems for evaluation. The remaining 527 problems have 16k valid solutions, from which we randomly sampled a subset for training. We format the text solutions as following format: >/24]", where (.../ represents current numbers, and (...) represents operations. The operations combine two numbers with an operator, resulting in new number. We train D-TSN to select operands and operators, and use the transition function to handle the computation of the current numbers after the operation, i.e. given state "/2.3,7.71->(2*** and action to select "7" as the next operand, the next state given by the transition function is "2.33.7|-((77414))337744444 would allow to focus more on the difficult part of the solution, which is deciding which numbers to use in the operation and which operator to use. Such transition can be easily obtained in practice we observe that the SFT model never makes mistakes these positions in its generation.
This

---

Published as a conference paper at ICLR 2025
4.1.3 RESULTS
We train D-TSN using search iterations in training and compare the resulting value function with supervised fine-tuned model, as described in Equation 17 to guide the search during evaluation.
Table I: Comparison between D-TSN and SFT on Game of 24, with varying amount of training trajectories, using metrics Success Rate whether the generated expression is valid and results in 24. D-TSN is trained and evaluated with 8 search iterations, SFT is evaluated with 8 search iterations.
#Training Trajectories
Method
1k
2k
10k
D-TSN (n itr=8)
31,46
38.25
45.24
SFT+Search
12.45
15.09
21.32
To compare D-TSN with a supervised fine-tuned model under different amounts of available data, we sample 1k, 2k, and 10k trajectories to train both methods. Table [ shows that D-TSN can better guide the search during evaluation when compared to the SFT model. On a closer look, we notice that the log- ligeliiod of the SFT model as value function often fails to recognize promising states, assigning a lower value to them, even when the states have already reached 24. As a result, tokens chosen at each step often do not lead to correct solution. In contrast, D-TSN better recognizes these states and assigns high values to them, guiding each step to select tokens that lead to a correct solution.
4.2 D-TSN WITH JOINTLY LEARNED WORLD MODEL
Next, we study two problems where we jointly learn world model in D-TSN. The states in these problems are represented as pixels of images.
Navigation This is 2D grid-based navigation task designed to quantitatively and qualitatively visualize the agent's s generalization capabilities. The environment is a 20 x 20 grid featuring a central hall. At the start of each episode, robot is randomly positioned inside this hall, while its destination is set outside. We present two distinct scenarios: one with a single exit and another with two exits from the central hall. Training is done only in the exit scenarios. The single exit scenario, which requires a longer-horizon planning to reach the goal, is used to test generalization.
Procgen Procgen (Cobbe et al., 2019) is a collection of 16 procedurally generated, game- environments, specifically designed to evaluate an agent's ' generalization capability, differentiating from Atari 2600 games (Mnih et al., 2013). Further details on these domains are presented in Appendix[..]
We follow the Offline-RL paradigm and train D-TSN using sequences of demonstrations. Please refer to Appendix D.2 for a detailed explanation of the learning framework.
4.2.1 TRAINING DATASETS
We use a behavior policy, which can be optimal or sub-optimal, to collect demonstration sequences for training. An optimal policy generates dataset with lower noise and a cleaner training signal, leading to a stable learning process. In contrast, sub-optimal policy produces a noisier dataset, which consequently restricts the quality of the policy that the agent can learn. The choice of the behavior policy depends on domain--seccinn requirements and the resources available for data collection. We use an optimal behavior policy for Navigation and sub-optimal policy for Procgen.
The training dataset, D, consists of trajectories generated using the behavior policy тB, where each trajectory defined as a series of T tuples, each comprising the state observed, action taken, reward observed, and the corresponding Q-value of the observed state, denoted as {(st, d;, ri, [ The Q-value for state 8; can computed by adding the rewards obtained in the trajectory from timestep onwards, i.e. Q = j=m² Tj- We limit the number of trajectories to 1000 for each domain to evaluate the sample complexity and generalization capabilities of each method.
4.2.2 BASELINES
We benchmark D-TSN against the following prominent baselines:
7

---

Published as a conference paper at ICLR 2025
Model-free Q-network This allows to assess the significance of incorporating inductive biases into the neural network architecture. This model is trained using the loss defined as:
Le L + X2LD
(18)
Model-based Search In this baseline, we utilize the submodules defined for D-TSN, but the world models and the value module are trained independently of each other. During evaluation, we employ the best-first search, akin D-TSN, utilizing the independently trained modules. Through this baseline, we assess the benefits derived from the joint optimization of the world model and the search algorithm. For our evaluations, we perform 10 search iterations for each input state. To train this model, we compute the Q-value, Qa. withoul performing the search and optimize the loss defined as:
Cearch X1 + AgCD + X3CT, + *rCR,
(19)
TreeQN This comparison helps highlight the advantages of advanced search algorithm, used in D-TSN, that can execute a deeper search while maintaining similar computational constraints. For evaluations, we adhere to a depth of 2 for TreeQN, as described in Farquhar et al. (2018), for both Procgen and navigation domains. Notably, the memory footprint would increase exponentially for larger depth, which prohibits to use larger depth. This model is discussed in Farquhar et al. (2018)) defined as:
trained using the loss (as
(20) Every method is trained with the same datasets using their respective loss functions. A more comprehensive discussion of the baselines and implementation details can be found in Appendix D.6
CTEEN = X1 + XgCD >gCR,
Table Comparison of D-TSN with the baselines on Procgen using metrics Mean Scores and Mean Z-score. Standard error intervals are reported, based on 1000 evaluation runs per game.
Table 2: Comparison of D-TSN with the baselines on Navigation using metrics Success Rate and Collision Rate.
Solver
Success Rate
Absolute Collision Rate
Games
Model-free
Q-netwrek
Modie--usee Search
TreeQN
D-TSN
bigfish
21.6710.53
16.3610.50
20.3810,52
21.76_00.53
Navigation exits)
bosslight
8.33+0.18
7.3340019
8.35-0.19
8,50=
cavetleer
2,05++0.12
3.52+0:15
3.57=0.15
Model-free Q-network
94.5% (# 0.29
4.4%
chaser
5.77::0.16
5.1100:15
6.37:0017
Model-based Search
93.2% (4 0.39)
6.7%
dlibeer
2.85+00115
4.33+0.17
4.13-0.16
5.62-0.15
TreeQN
95.4% (202
3.8%
councun
dudgebail
5,18+0.16
1,05100006
6.33++0.15
4.26100177
5.01=0.16
4,7020.18
5,30 0.19
6.40=0.15
D-TSN
99.0%
(2
0.7%
fniiforr
14.09++0.57
10.76-00.8
15.43+0.35
14:00-0.37
2.24110.3
1.1110.12
Navigation (1 exit)
jumper
3.21
3.97100155
3.7420.15
4.65 0.16
Model-free Q-network
47.1% (+:55))
50.2%
maze
lesper
7.78++0.13
2.0210.13
2.2611..1
6.3-+0.15
2.5010..4
7.85-0.13
8.29- 0.12
1.30
Model-based Search
86.9% (t 0.59)
12.4%
141110.02
1.3510.05
2.0610.00
1.71 10.6
TreeQN D-TSN
51.8 058)
39.2%
ninja
5.07+0.16
5.0750.16
5.05
6.61=0.15
99.3% (1 (1.)
0.2%
plunder
13,7710.36
9.7410.32
12.75_0335
sturpilot
15.7410..99
14.3110.88
17.9110.42
16.12_0041
Meon Z-Seww
0.10
0.27
0.35
4.2.3 RESULTS
Navigation We compare D-TSN against the baselines using success rate and collision rate as evaluation metries, where success rate refers to the fraction of test levels completed by the agent, and collision rate refers to the fraction levels failed due to collision with a wall. Results are shown in Table 2
We observe that D-TSN outperforms the baselines on both navigation scenarios. Notably, when agents trained on data from the 2-exit scenarios but tested in the 1-exit scenario. D-TSN, with its powerful inductive bias, retains its performance. In stark contrast, the Model-free Q-network and TreeQN experience substantial performance decline, which underscores their limited generalization ability. Model-bo based Search also registers a minor decrease in the success rate, reinforcing the importance of jointly optimizing the world model for enhanced robustness.
Procgen We evaluate the performance of D-TSN and the baselines on all Procgen suite games, comparing mean Z-score and head-to-head wins. We also report the mean scores across 1000 episodes obtained by these methods in Table 3 As Procgen games have different scales, we use model-free
Q-network as baseline compute a normalized score Z Z-score
(μ* /BBB / B, where fx and μy
We use 7-score because the baseline normalized score used in prior works Badia al. 2020) Kaiser et al.
2 Mittal al. 2023) is problematic our experiments, as described in Appendix D.7
8

---

Published as a conference paper at ICLR 2025
represent the mean scores obtained by the agent policy and the baseline policy respectively and øg represents the standard deviation of the scores obtained by the baseline policy.
It is important to note that the training data for Proegen generated using sub-optimal behavior policy, leading to noisy training signals in comparison to the relatively noise-free data in the Naviga- tion domain. Despite this, D-TSN reports higher mean Z-score, averaged across the 16 Procgen games, particularly in climber, coinrun, jumper, and ninja, which require long-term planning. This underscores the stronger inductive bias of D-TSN. In Procgen, Model- Search lags behind both TreeQN and D-TSN, highlighting that for complex environments, joint optimization enables learning robust world model. In the head-to-head comparison listed in Table D-TSN won in 13 games against the Model-free Q-network, 16 games against Model-bart -based Search, and 13 games against TreeQN.
Table 5: Comparison of D-TSN with its variants to evaluate the role of Auxiliary losses, REIN- FORCE term and Telescoping Sum.
Table 4: Head-to-hee heah comparison of D-TSN with the baselines on Proegen using metric number of games won.
Navigation (1-exit)
Tree Search Network
Solver
Proegen
Baseline
(Games won Total games)
Success Rate
Mean Z-score
Model-free Q-network Model-based Search
13 / 16
D-TSN
99.3% (d: 0.1%)
16 / 16
w/o Telescoping Sum
98.5%
0.28
TreeQN
13 / 16
w/o REINFORCE term
97.7%
(t: 0.2%)
0.29
wlo Auxiliary losses
91.1% (#: 0.3%)
4.3
ABLATION STUDIES
4.3.1 THE IMPACT OF THE REINFORCE TERM AND TELESCOPING SUM TRICK
In this study, we assess the impact of both the REINFORCE term and the Telescoping Sum trick on D-TSN's performance. The results, presented in Table show notable differences. Without the telescoping sum, D-TSN sees a modest decrease in the success rate for Navigation (1-exit), moving from 99.3% 98.5%. Similarly, the mean Z-score for Proegen dips to 0.28. The omission of the REINFORCE term also marks a decline, with the Navigation (1-exit) success rate landing at 97.7% and the mean Z-score for Procgen dipping to 0.29.
4.3.2 THE CONTRIBUTION OF AUXILIARY LOSSES
In this study, we explore the contribution of auxiliary losses to D-TSN's performance. As outlined in T5 the absence of these auxiliary losses leads more pronounced decline in the performance. Specifically, the success rate in the Navigation (1-exit) task drops significantly to 91.1 % compared to 99.3% with auxiliary losses.
4.3.3 ENHANCING PERFORMANCE THROUGH DEEPER SEARCH
In this study, explore the benefits of deeper search in the D-TSN framework by increasing search iterations in Navigation and Game of 24 problems. For Navigation, we maintain an identical number of iterations in the training and evaluation phases to prevent any distribution shifts in the world model. For Game of 24, we can use a larger search iteration during evaluation, as the world model remains fixed.
Table 6: Comparison of D-TSN trained with different number of search iterations ('n_itr') to evaluate the performance gain by perform- ing deeper searches.
Table Comparison of D-TSN using different num- ber of search iterations during evaluation. All methods are trained with 2k trajectories.
#Scarch iterations in evaluation
Solver
Navigation (2-exits)
Navigation (1-exit)
16
128
D-TSN (n iit=l)
20.75
31.51
38.16
49.76
89.29
D-TSN (n_ itr=S)
97.4%
96.6%
D-TSN (n itr=))
32.22
41.13
49.67
90.24
D-TSN (n itr=10)
99.0%
99.3%
D-TSN (n 8)
38.25
49.91
91.42
D-TSN (n 20r
99.5%
99.1%
D-TSN (n itr=
49.58
92.38
9

---

Published as a conference paper at ICLR 2025
The results, as listed in Table 6 and Table 7 that a higher number of iterations leads to an improvement in the success rate in both Navigation and Game of 24. For Navigation, where the world model is not known, the improvement suggests that D-TSN is able to exploit the jointly learned world model effectively, and can potentially be scaled up further given additional computational resources. For Game of 24, see performance consistently improve the number of search iterations during evaluation increases, highlighting the importance of search. Notably, the performance reaches around 90% for D-TSN when there are sufficient search iterations during evaluation, suggesting that the value function learned by D-TSN can recognize the promising states and guide the generation to those states. We also observe that the number of search iterations performed during training does not affect the performance much in Game of 24. One explanation is that we are using a ground-truth world model, and performing a search using actual states. The REINFORCE term Equation (IT) would act similarly to the CQL loss in Equation (13). Learning world model jointly in D-TSN is also possible in principle in language model tasks, but scaling considerations become important to manage the large action space and observation space. We leave this exploration for future works.
4.3.4 ROBUSTNESS OF THE WORLD MODEL.
A key challenge in employing learned world models for search is addressing the compounding errors, which impact the accuracy and effectiveness of search. This study shows that the joint optimization of both the world model and the search algorithm compensates for these inaccuracies, ensuring the world model usable in deeper online searches.
Table 8: Comparison of D-TSN and Model- Search to highlight the robustness of the world model when performing deeper searches ('n ('n_itr' refers to number of search iterations).
Navigation (2-exits)
Navigation (1-exir)
Solver
D-TSN (n_itr=10) 10) D-TSN (it=20))
99. 0%%
99.3%
99.3%
99.49.
D-TSN (n_itr (n_itr-50)
99.7%
99.6%
Model-based Search (n_ it)=10))
93.2%
86.9%
Model-based Search (it (0=
91.1%
84.6%
Model-based Search (n. (tm
89.5%
80.4%
As demonstrated in Table 8, the world model trained jointly with the search algorithm consistently outperforms the independently trained model, especially as the number of iterations increases.
CONCLUSION
In this paper, we introduce Differentiable Tree Search Network (D-TSN), novel framework that learns to conduct differentiable tree search from just sequences of demonstrations. D-TSN can leverage known world model, or jointly learn world model end-to-end with other search submod- ules. D-TSN conducts a best-first style search, which allows to search deeper and explore more promising states. A naive incorporation of best-first search could lead to discontinuity of the loss in the parameter space, an issue we address by employing a stochastic tree expansion policy. We optimize the expected loss function using REINFORCE-styL objective and propose telescoping sum trick to reduce the variance of the gradient for this expected loss. We evaluate D-TSN in two scenarios: when the world model known we experiment with language model reasoning task, Game of 24: 2) when the world model is not known we experiment with 2D grid navigation and Proegen games. Through our experiments, we show that D-TSN effective, outperforming baselines in both scenarios, especially when the world model is jointly learned.
Nonetheless, the strength of the current implementation of D-TSN is currently limited to deterministic decision-making problems with discrete action space. We have demonstrated the method on latent world models with 2D visual observations but have yet to do so on language observations. To cater to a broader spectrum of ceciisionmmainng problems, there is also need to revamp the transition model to manage stochastic world scenarios, which include common usage of language and visual language models. We aim to address these problems in our future works.
ACKNOWLEDGMENTS
This research is supported by the Ministry of Education, Singapore, under its Academic Research Fund Tier 1 (A-80008114000000)
10

---

Published as a conference paper at ICLR 2025
REFERENCES
Kavosh Asadi, Dipendra Misra, and Michael L. Littman. Lipschitz continuity in model-based reinforcement learning. In Jennifer G. Dy and Andreas Krause (eds.), Proceedings of the 35th International Conference on Machine Learning, ICMI. 2018, Stockholmsmüssann Stockholm, Sweden, July 10-15, 2018, volume 80 of Proceedings of Machine Learning Research, pp. 262-26442777 273. PMLR, 2018, URL
html
Adrià Puigdomènech Badia, Bilal Piot, Steven Kapturowski, Pablo Sprechmann, Alex Vitvitskyi, Zhaohan Daniel Guo, and Charles Blundell. Agent57: Outperforming the atari human benchmark.
In Proceedings of the 37th International Conference on Machine Learning. JCML 2020. 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 507-517. PMLR R, 2020. URL httpt:::p/ppoceeednnge
m prwpsressssee /vlas//vlii wwtpss/wwee//eeeeeeeeeeeeeeeeeee 19/bad1a20aaae att.mmmmm
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. arXiv preprint arXiv: 1912.01588, 2019.
Karl Cobbe, Christopher Hesse, Jacob Hilton, and John Schulman. Leveraging procedural generation to benchmark reinforcement learning. In Proceedings of the 37th International Conference on Machine Learning, ICML 2020, 13-18 July 2020, Virtual Event, volume 119 of Proceedings of Machine Learning Research, pp. 2048-2056. PMLR, 2020. URL http: http:::/prcceeeinngg
mlr wtw.press/vll9.. 1/ http:///1.///bbbbblbbaa.. l
Karl Cobbe, Jacob Hilton, Oleg Klimov, and John Schulman. Phasic policy gradient. Proceed- ings of the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, 339:2020-2021, 2021. URL http://proceediing www.smmrr..... wtt.:gs.......p..s113311
cobbe21a. html
Fernando Fradique Duarte, Nuno Lau, Artur Pereira, and Luis Paulo Reis. A survey of planning and learning in games. Applied Sciences, 10(13):4529, 2020.
Abhimanyu Dubey, Abhinav Jauhri, Abhinav Pandey, Abhishek Kadian, Ahmad Al-Dahle, Aiesha Letman, Akhil Mathur, Alan Schelten, Amy Yang, Angela Fan, Anirudh Goyal, Anthony Hartshom, Aobo Yang, Archi Mitra, Archie Sravankumar, Artem Korenev, Arthur Hinsvark, Arun Rao, Aston Zhang, Aurélien Rodriguez, Austen Gregerson, Ava Spataru, Baptiste Rozière, Bethany Biron, Binh Tang, Bobbie Chern, Charlotte Caucheteux, Chaya Nayak, Chloe Bi, Chris Marra, Chris McConnell, Christian Keller, Christophe Touret, Chunyang Wu, Corinne Wong, Cristian Canton Ferrer, Cyrus Nikolaidis, Damien Allonsius, Daniel Song, Danielle Pintz, Danny Livshits, David Esiobu, Dhruv Choudhary, Dhruv Mahajan, Diego Garcia-Olane Diego Perino. Dieuwke Hupkes, Egor Lakomkin, Ehab AlBudawy, Elina Lobanova, Emily Dinan, Eric Michael Smith, Filip Radenovic, Frank Zhang, Gabriel Synnaeve, Gabrielle Lee, Georgia Lewis Anderson, Graeme Nail, Grégoire Mialon, Guan Pang. Guillem Cucurell, Hailey Nguyen, Hannah Korevaar, Hu Xu, Hugo Touvron, Iliyan Zarov, Imanol Arrieta Ibarra, Isabel M. Kloumann, Ishan Misra, Ivan Evtimov, Jade Copet, Jaewon Lee, Jan Geffert, Jana Vranes, Jason Park, Jay Mahadcokar, Jeet Shah, Jelmer van der Linde, Jennifer Billock, Jenny Hong, Jenya Lee, Jeremy Fu, Jianfeng Chi, Jianyu Huang, Jiawen Liu, Jie Wang. Jiecao Yu, Joanna Bitton, Joe Spisak, Jongsoo Park, Joseph Rocca, Joshua Johnstun, Joshua Saxe, Junteng Jía, Kalyan Vasuden Alwala, Kartikeya Upasani, Kate Plawiak, Ke Li, Kenneth Heafield, Kevin Stone, and et al. The llama 3 herd of models. CoRR, /tbb//240.221788 2024. doi: 1048550/ARXIV.2407.2777 48550 2ttp555/axxiv.240..... .17833
URL https
Gregory Farquhar, Tim Rocktäschel, Maximilian Igl, and Shimon Whiteson. Treeqn and atreec: Differentiable tree-structured models for deep reinforcement learning. In 6th International Confer- on Learning Representations, ICLR 2018, Vancouver, BC, Canada, April 30 May 3, 2018, Conference Track Proceedings. OpenReview.net, 2018. URLhhttp:://ppeeeeieee net/
Eorum d= d= d
Arthur Guez, Théophane Weber, Ioannis Antonoglou, Karen Simonyan, Oriol Vinyals, Daan Wierstra, Rémi Munos, and David Silver. Learning to search with metsnets. In International Conference on Machine Learning, pp. 1822-1831. PMLR, 2018

---

Published as a conference paper at ICLR 2025
David Ha and Jürgen Schmidhuber. Recurrent world models facilitate policy evolution. In Ad- vances in Neural Information Processing Systems 31: Annual Conference on Neural Informa- fion Processing Systems 2018. NeurlPS 2018, December 3-8, 2018, Montréal, Canada, pp. 2455-2467, 2018. URL /tw.tts::pproceeeeeeeeeee neur www.errrs..c...eoo ttt.cccppppppp wwtpcc//pppp.....c 2atar///0//hhhhhh 8/hash/ 7865016762 -bw.::/22aaaaaaaaaaaaaaaaaaaa-a- html
Danijar Hafner, Timothy P. Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. In 8th International Conference on Learning Repre- sentations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net 2020. URL
net/
Shibo Hao, Yi Gu, Haodi Ma, Joshua Jiahua Hong, Zhen Wang, Daisy Zhe Wang, and Zhiting Hu. Reasoning with language model is planning with world model. In Houda Bouamor, Juan Pino, and Kalika Bali (eds.), Proceedings of the 2023 Conference on Empirical Methods in Natural Language Processing, EMNLP 2023, Singapore, December 6-10, 2023, pp. 8154-811 173. Association for Computational Linguistics, 2023. doi: 10. (0.186633/V122223 EMW.:023...MMM--MMMMAMM MAIN.5007.. www.222222mmmmmmmmmmmmmmmmmmm wwweemnnnmemee in. 507
URL httppp///oool.oo//..0.66333 215.:/10.166656
Lukasz Kaiser, Mohammad Babacizadch, Piotr Milos, Blazej Osinski, Roy H. Campbell, Konrad Czechowski, Dumitru Erhan, Chelsea Finn, Piotr Kozakowski, Sergey Levine, Afroz Mohiuddin, Ryan Sepassi, George Tucker, and Henryk Michalewski. Model based reinforcement learning for atari. In 8th International Conference on Learming Representations, ICLR 2020, Addis Ababa, Ethiopia, April 26-30, 2020. OpenReview.net, 2020. URL https:/ https::/openrevieew openrevien.eee.eeeoe net/
Corum?id-Slxcn JH.
Aviral Kumar, Aurick Zhou, George Tucker, and Sergey Levine. Conservative q-learning for offline reinforcement learning. Advances in Neural Information Processing Systems 33: Annual Conference Neural Information Processing Systems 2020, NeurIPS 2020, December 6-12, 2020, virtual, 2020. URL /ttp: 0d2 62061 wwt.222222ff2222a2 6a50 Abs. html
Sergey Levine, Aviral Kumar, George Tucker, and Justin Fu. Offline reinforcement learning: Tutorial, review, and perspectives on open problems. CoRR, abs/2005.01643 2020. URL https://
613
Jiacheng Liu, Andrew Cohen, Ramakanth Pasunuru, Yejin Choi, Hannaneh Hajishirzi, and Asli Celikyilmaz. Don't 't throw away your value model! generating more preferable text with value- guided monte-carlo tree search decoding, 2023.
Siqi Liu, Kee Yuan Ngiam, and Mengling Feng. Deep reinforcement learning for clinical decision support: A brief survey. CoRR, 00s//997.00945, 2019. URL httt:::/arxiv.ogg// 1907.000 09475
Dixant Mittal, Siddharth Aravindan, and Wee Sun Lee. Expose: Combining state-based exploration with gradient- based online search. In Proceedings of the 2023 International Conference on Autonomous Agents and Multiagent Systems, AAMAS '23, pp. 1345-1353, Richland, SC, 2023. International Foundation for Autonomous Agents and Multiagent Systems. ISBN 9781 150394333.
Volodymyr Mnih, Koray Kavukcuoglu, David Silver, Alex Graves, Ioannis Antonoglou, Daan Wierstra, and Martin A. Riedmiller. Playing atari with deep reinforcement learning. CoRR, abs/1312.5502,, 2013. /tLp://t://aat
org/ abs /t///1111... 5602
Volodymyr Mnih, Adrià Puigdomènech Badia, Mehdi Mirza, Alex Graves, Timothy P. Lillicrap, Tim Harley, David Silver, and Koray Kavukcuoglu. Asynchronous methods for deep reinforcement learning. In Maria-Florina Balcan and Kilian Q. Weinberger (eds.). Proceedings of the 33nd International Conference on Machine Learning, ICML 2016, New York City, NY, USA, June 19-24, 2016, volume 48 of JMLR Workshop and Conference Proceedings, pp. 1928-1937, JMLR.org, 2016. URL hta.lttt::// httt:::/proceed:: ing...g.g.... m www.ppresssse. v 8/ mm ha 6.htm
M. G. Mohanan and Ambuja Salgoankar. A survey of robotic motion planning in dynamic environ- ments. Robotics Auton. Syst. 100:171-188 2018 doi: 10.1016/J R.0.6//...B...001.... 10.001. URL
https://doo .tt./oorgo.... wttp///..11111//11. 10w.10066.... robot 2017 10.011
12

---

Published as a conference paper at ICLR 2025
Soroush Nasiriany, Vitchyr Pong. Steven Lin, and Sergey Levine. Planning with goal- conditioned policies. In Hanna M. Wallach, Hugo Larochelle, Alina Beygelzimer, Flo- rence d' d'Aché-Buuc, Emily B. Fox, and Roman Garnett (eds.), Advances in Neural In- formation Processing Systems 32: Annual Conference on Neural Information Process- ing Systems 2019, NeurlPS 2019, December 8-14, 2019, Vancouver, BC, Canada, pp. 14814-14825, 2019, URL wtLLitttps//ppoceeeiiiii http://oocccennnn....iu....... wttps/urrs.....ppppeeeeee/ee
Abst html
Eduardo Pignatelli, Johan Ferret, Matthieu Geist, Thomas Mesnard, Hado van Hasselt, and Laura Toni. survey of temporal credit assigmment in deep reinforcement learning. Trans. Mach. Learn. Res., 2024, 2024. URLptttss: htt.ttss://oeerev..eee www.:/eettoornieeo forum?irummmm d=bNt
net
Rafael Figueiredo Prudencio, Marcos Ricardo Omena Albuquerque Máximo, and Esther Luna Colombini. A survey on offline reinforcement learning: Taxonomy, review, and open problems. IEEE Trans. Neural Networks Learn. Syst., 35(8): 2024. doi: 10. 2033. 3250269. URL htthh:ht/ooo..... org 00. www.o/////.....//...........11.e 1111/
Walter Rudin. Principles of Mathematical Analysis. McGraw-Hill New York, 3d ed. edition, 1976. ISBN 007054235.
html
Julian Schrittwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Simon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with learned model. Nature, 588(7839):604-669, 2020.
John Schulman, Nicolas Heess, Theophane Weber, and Pieter Abbeel. Gradient estimation using stochastic computation graphs. In Corinna Cortes, Neil D. Lawrence, Daniel D. Lee, Masashi Sugiyama, and Roman Garnett (eds.), Advances in Neural Information Processing Systems 28: Annual Conference on Neural Information Processing Systems 2015, December 7-12, 2015, Mon- treal, Quebec, Canada, pp. 3528 3536, 2015. httpsunntttpp:////eeeeeeiiieeeo neur ips. cc/ .html
Max Schwarzer, Ankesh Anand, Rishab Goel, R. Devon Hjelm, Aaron C. Courville, and Philip Bachman. Data-efficient reinforcement learning with self-predictive representations. In 9th Interna- fional Conference Learning Representationss ICLR 2021, Viriual Event, Austria, May 3-7, 202 OpenReview.net, 2021. URL / wwte/foouue forumm wtt.mmmmiiduuuuguuuUu.u
Ameesh Shah, Eric Zhan, Jennifer J. Sun, Abhinav Verma, Yisong Yue, and Swarat Chaud- huri. Learning differentiable programs with admissible neural heuristics. In Hugo Larochelle, Mare' Aurelio Ranzato, Raia Hadsell, Maria-Florina Balcan, and Hsuan-Ties Lin (eds.), Advances in Neural Information Processing Systems 33: Annual Confer- ence Neural Information Processing Systems 2020, NeurlPS 2020, December 6-12, 2020, virtual, 2020, URL /tL/itttpss//pppceeeeiiee neurips. cc 2aaper/2020/hahl/ Abstract html
David Silver, Aja Huang, Chris J. Maddison. Arthur Guez, Laurent Sifre, George van den Driess- che, Julian Schrittwisser, Ioannis Antonoglou, Vedavyas Panneershelvam, Marc Lanctot, Sander Dieleman, Dominik Grewe, John Nham, Nal Kalchbrenner, Ilya Sutskever, Timothy P. Lillicrap, Madeleine Leach, Koray Kavukeuoglu, Thore Graepel, and Demis Hassabis. Mastering the game of go with deep neural networks and tree search. Nature, 529(7587):484 2016. doi: 10.10088/aaurr16660 16961. URL httph:ttt//dttt
110 10300 6961
Bharat Singh, Rajesh Kumar, and Vinay Pratap Singh. Reinforcement learning in robotic ap- plications: a comprehensive survey. Artif. Intell. Rev. 55(2): 990, 2022. doi: 10.1007/ S10462-021-09997..9 URL https:://ds/
100.//010.. 1007 0010://0000000100420000001111110
Adish Singla, Anna N. Rafferty, Goran Radanovic, and Neil T. Hefferman. Reinforcement learning for education: Opportumities and challenges. CoRR, abs/2107.08822 2021.
URL https
Richard S Sutton and Andrew G Barto. Reinforcement learning: An introduction. MIT press, 2018.
13

---

Published as a conference paper at ICLR 2025
Ronald J Williams. Simple statistical follown algorithms for connectionist reinforcement learning. Machine learning, 8(3):229-2) 256, 1992.
Yuxi Xie, Kenji Kawaguchi, Yiran Zhao, James Xu Zhao, Min-Yen Kan, Junxian He, and Michael Qizhe Xie. Self-evaluation guided beam search for reasoning. In Alice Oh, Tris- tan Naumann, Amir Globerson, Kate Saenko, Moritz Hardt, and Sergey Levine (eds.), Ad- vances in Neural Information Processing Systems 36: Annual Conference on Neural Infor- mation Processing Systems 2023, NeurlPS 2023, New Orleans, LA, USA. December 10 2023, 2023. 11a.555h6....33ahh22222263216610 81 f1
Shunyu Yao, Dian Yu, Jeffrey Zhao, Izhak Shafran, Thomas L. Griffiths, Yuan Cao. and Karthik Narasimhan. Tree of thoughts: Deliberate problem solving with large language models. CoRR, abs/2305. 2023. doi: 40.48555//RXXI.23055. 48550/arxiv 10601
URL htt.t:::////oo.rrtttt. org/
Weirui Ye. Shaohuai Liu, Thanard Kurutach, Pieter Abbeel, and Yang Gao. Mastering atari games with limited data. In Marc' Aurelio Ranzato, Alina Beygelzimer, Yann N. Dauphin, Percy Liang, and Jennifer Wortman Vaughan (eds.), Advances in Neural Information Processing Systems 34: Annual Conference on Neural Information Processing Systems 2021, NeurlPS 2021, December 6-14, 2021, virtual, pp. 25488825448, 2021. URL http:tttss::ppcceeeeeeeeeee neurips.. cc/ wwweerbtttttttreeeeeeeeeeeeeee.e hhmmll
Ryo Yonetani, Tatsunori Taniai, Mohammadamin Barekatain, Mai Nishimura, and Asako Kanezaki. Path planning using neural a* search. In Marina Meila and Tong Zhang (eds.). Proceedings the 38th International Conference on Machine Learning, ICML 2021, 18-24 July 2021, Virtual Event, volume 139 of Proceedings of Machine Learning Research, pp. 12029-12039, PMLR, 2021. URL http://proceedingg www.innnn..mm....mm ptte: ltt
Andy Zhou, Kai Yan, Michal Shlapentokh-Rotthmal Haohan Wang, and Yu-Xio Yu-Xiong Wang. Language agent tree search unifies reasoning, acting, and planning in language models. Forty-first International Conference on Machine Learning, ICML 2024, Vienna, Austria, July 21-27, 2024. OpenReview.net, 2024. URL https: /os.ss.ss...ttt.///ee/e.rro..o.. nww. www.:orwwmmmdddnn......m..
14

---

Published as a conference paper at ICLR 2025
A DIFFERENTIABLE TREE SEARCH NETWORK ALGORITHM
A.1
DIFFERENTIABLE TREE SEARCH NETWORK PSEUDO-CODE
Algorithm 1: Differentiable Tree Search (D-TSN)
Input: Input state, 8root
Result: Q-values, Q(Sroot: Q)
hroot & Eg (sroot)
//Encode so to its latent state ho //Initialize root node //Initialize candidate set Open
noderoot t initialise(Prool)
Open {noderoot):
// Expansion phase
itr 0-
repeat
foreach node E Open do
hnode + (etdattntnnode)
//Get latent state of node (node) + sumO Rewards(node)) + Vo(noode): / //Compute path values
Teree + softmax,, (V(n)):
/ //Compute the tree expansion policy //Sample the node to expand
node" sample (Ttree)
foreach a E Actions; do
hehild +- Tonodee a)
/ //Compute the next latent state / //Compute reward for the transition //Create the child node
Fehild < (-Roo (hnoer a):
create.Node CheNoodeehhhid.... Fehita):
Open + + Open {child, cchid, = getChild(node" .a) E A} - node"; //Update
the open set
ilr ilr + 1; until ilr < MAX_ITR; Backup phase
foreach node E Tree, iterating from leaf nodes the root node; do
node leaf: then
Anode + gel Laten((n(dd)
//Get latent state of node //Compute value using Value module
((noee) = (hnooed)
else
foreach a E Actions; do nodechsidaa + gelChiildnnode;))
/ //Get child of node that
corresponds to action
Bnode + Latentl((nod)
//Get latent state of node //Get reward using Reward module
Fnode +- Ro (h a)
Qnoode, a) Noode + V (nodecnsdja))
(node) + maxa Q(node,a) a)
return ((ooeerott):
//Return Q-value of the root node
A.2 PSEUDO-CODE EXPLANATION
D-TSN initializes the search by converting the input state its latent representation and proceeds in two phases: Expansion and Backup as detailed in the following sections.
A.2.1 INITIALIZATION
Given an input state Syoots the algorithm begins by encoding this state into its latent representation Tr using an encoder Eg. This latent representation serves as the root node of the search tree.
15

---

Published as a conference paper at ICLR 2025
A.2.2 EXPANSION PHASE
The algorithm initiates set of candidate nodes, termed 'Open', starting with the root node.
In each iteration, the algorithm considers every node in the 'Open' set, retrieves its latent state, and computes an interim value (node), which combines the cunulative rewards of the node's path with value estimation from the value module Vo-
Using the path values of nodes in 'Open', the tree expansion policy, Tuee. is computed. From this policy, node, node", is sampled for expansion.
Each possible action from node* results in the creation child node. This achieved by leveraging the transition module To to predict the latent state of the child and the reward module R. to determine the associated reward. After expansion, node. * is removed from 'Open' and its children are added.
This expansion process continues until a predetermined number of iterations, MAX_ LTR, is reached.
A.2.3
BACKUP PHASE
Starting from the leaf nodes, the algorithm backpropagates value estimates to the root.
For each leaf node, its value is directly computed from the value module Vg. For non-leaf nodes, the Q-value for each action is estimated by combining the reward for that action with the value of the corresponding child node.
The value of non-leaf node is set to the maximum Q-value among its actions.
A.2.4 OUTPUT
The algorithm finally returns the Q-values associated with the root node, ((nodeoa)) providing estimation of the value of taking each action from the initial state.
This explanation provides high-level view of the D-TSN algorithm's operation. By breaking down the search process into expansion and backup phases, the pseudo-code highlights how D-TSN incrementally builds the search tree and then consolidates value estimates back to the root.
B CONTINUITY OF THE LOSS FUNCTION
In this section, we show that the Q-function represented by a search tree with fixed structure is continuous network's parameter space.
Suppose we have two functions fz) and g(z) which are continuous at any point c in their domains. Lemma B.1. Continuity of Composition: The composition of two continuous functions, denoted as f((())), retains continuity. (Theorem 4.7 in Rudin (1776))
Lemma B.2. Continuity of Sum operation: The result of adding two continuuss expressed as fx) ((x), is continuous function. (Theorem 4.9 in Rudin 19761)
Lemma B.3. Continuity of Max operation: Applying Max as max(f(x).g(() results a function that is continuous.
continuous functions, expressed
Proof. Consider a function h(x) = (ax( g(x)), and we aim to demonstrate that h(x)) is continuous. Now, hxz) can be expressed as combination of continuous functions:
h(x)) f(x) + (xx) + [f(x) -(x)
2
Since sums and absolute values of continuous functions are continuous (Rudin 1976). h(x)) is continuous. Hence, the maximum of two continuous functions is also continuous function.
Rewriting the Theorem 3.1 with proof below:
Theorem B.4. Given a set of parameterised modules that are continuous in the parameter space 8. the Q-function computed by expanding fixed search tree by composing these modules and backpropagating the children values using addition and max operations continuous in the parameter space O. When the tree structure is not fixed, the continuity of the Q-values not guaranieed.
16

---

Published as a conference paper at ICLR 2025
Proof. Let represent the set of parameters, encoder module, transition module, reward module, and value module respectively as 8, Eg. Tg. Rg, and Vg. These submodules are assumed to be composed of simple neural network architectures, comprising linear and convolutional layers, with the Rectified Linear Unit (ReL.U) serving as the activation function. These submodules, therefore, are continuous within the parameter space.
We can subsequently rewrite the Q-value as the output of a full tree expansion as follows:
Q(so, ao) = Q( ao) = ro + V(h1)
(21)
where,
ho Eg(so)
(22)
Tt Ro(h, Uz)
(23)
h1+1 To( ax)
(24)
Qhe ax) = Ro az) + V(h++1)
(25)
(h+) VV (he)
if he a leaf otherwise
(26)
max, (Q(h,
Given that Ea. Ro, and To are continuous, ho, rt, and h441 in Equations continuous (derived from Lemma B.1.
(2 and 24 are similarly
Further, V(hg) in Equation 226 can either be (he),, if he is a leaf node, or max (Q a)) otherwise. In the first scenario, (h,, retains continuity by assumption. In the second scenario, if Q(he,as) (z) continuous, then (he) remains continuous as per Lemma B.3
Now, we show that Q(h,, ar) is continuous using recursive argument that for any node in the search tree, if the Q-values of all its child nodes are continuous, then Q-value is also continuous. Q-value of an internal tree node h, can be written Q( az) R& ae) + ((b++1), where h&+1 = The at) is the child node of ht. Considering the base case, when the hr41 is leaf node, then Q(h,, ..) Ro(hs, a.) + Vo (h,+2), which is continuous as per Lemma B.2 Consequently, (h+) maintains continuity. Applying this logic recursively, the Q-value Q Q(h,, 0z) of all the tree nodes maintains continuity.
Thus, we can decompose the Q-value at the root node, Q(so, 4n), as a composition of continuous functions, ensuring that the output Q-value, Q(5o, an), continuous.
Next, we show that the continuity of Q-function is not guaranteed when the structure of the search tree is not fixed by giving an example.
Consider a simple environment with two possible actions an and d2, and a state is represented by 2-D vector. We parameterize the reward, value, transition functions simple matrices: R2 R² is reward function that maps a state to a real number representing the reward for doing action a, the state; V 2² is the value function that maps state real number representing the value of a state; T; E _2x2 is a transition function that maps a state to a next state after doing action a; Consider search tree with root node E IR² and a child node x E RR2 derived by doing has two child node nı and n2. The path value of nı, defined by Equation (l. is V(n) = Ro(s, a1) + Re (I. a1) + Vo(nn) = R2 + R2 x + VT Similarly, the path value of R2 x + VTT..
action
N2 is ((n2) = (s( + Ro (I, a2) + (o R4 8
Consider we choose between ni and n2 for the next expansion of the tree. When R x: + VIT > RÍ x + VTT2,, n, is expanded, two children of n, has value R T + VTA and R Ti VTITTa, value of æ after backup is the maximum of value of mg and the values of two children of N1, which is max R2 x + VTTzg, max (R] Txx + VT R Tix TTT,,, }}.
When RT x VJ < RT VTT2r, ny is expanded, two children of n2 has value R T2x + TTT, and R Txx + TTTTT, value of after backup is maximum of value of n, and the values of two children of n2, which is max(R R x + VTT, x. max R Txx + VTT, T2x, R Txx VTT, }}.
17

---

Published as a conference paper at ICLR 2025
KCP
F33
F22
731
T:T
F24
T2f
00' > 2
2
Figure Visualization of discontinuous example.
Let's assign the following values to the mentioned symbols:
2 2
0
0 1 2
1
V
R₂
T
T2
We have R. VTTV' 40 and R x Thx 20 8. When 0 > n1 will be expanded, and when < 2, n2 will be expanded, see Figure 2 for visualization of the example. We have
lim (1)
= RT + lim max(R2 + TT,,, max(R + TTT,T R Th + (2TT,,) (27) a
=16
lim ((s,a1)
0-22-
= =R. + lim max( x + VT Tix max( T2x + T T/T2,, R T2x + T2T22,,) (28)
6-2
=32
Combining Equation 27 and Equation
28 we have that Q(r, a1) is discontinuous at 0 = 2.
C DERIVATION OF THE GRADIENT OF THE EXPECTED LOSS FUNCTION
Let us represent the Q-values predicted by D-TSN as Qo a/7), which depends on the final sampled after T trials of the online search. Let us denote the corresponding loss function on this output Q-value as C (0(ssa/T)) Our objective is to compute the gradient of the expected loss value, averaging over trees sampled.
18

---

Published as a conference paper at ICLR 2025
The gradient of expected loss, considering the expectation over the sampled trees, is derived as:
Vol(s,a)= - Vol., [((200.017))
= Ve Σ no To(T) ((0(s(((0())
Σ vo vo [T0())((Q0(e.att)
Σ [ VOTOTT) + xo(r) [=((t)Vo((oo(ssdt
) (T)) Σ [o(())oc(Ro((sot)))
= E (((0((, lolgma Vo((o[(..a))
I II mo(cul))+
E. (((0(s.a17) Vo log
= B, (((0(()
V log Toll
Leveraging the telescoping sum trick, as elaborated in Section 3.7 can be rewritten lower-variance estimate:
the gradient of the expected loss
Ve log To (nclT))) Voc((o(s,
VoL(s, = E,
where
R₂ Σr: ri rr= LT fL-1
f. 4= ((0(s, representing the loss value after the hth search iteration.
In practice, utilize the single-sample estimate for the expected gradient, elaborated in Schulman et al. 2015)
D
EXPERIMENTS
D.1 TEST DOMAINS
We grid navigation task and the Procgen games study D-TSN under the scenario when world models are not available. A visualization of these two tasks can be viewed in Figure
D.1.1 NAVIGATION
The grid navigation task serves as foundational that mimics the challenges robot might face when navigating in a 2D grid environment. This environment provides both a quantitative metric and qualitative visualization understand an agent's capacity to generalize its policy. Specifically, this task involves 20 x 20 grid with central hall. At the beginning of each episode, the robot is positioned at a random point within this central hall. Simultaneously, goal position sampled randomly location outside the hall, challenging the robot to find its way out and reach this target.
There are two variations of this task. The first provides the robot with a single exit from the central hall, while the second offers two exits. The single-exit hall scenario is similar to the two-exit scenario but requires a longer- horizon planning to successfully evade the walls to exit the hall and reach the goal.
19

---

Published as a conference paper at ICLR 2025
Navigation 2 exits
Proagen Games
Navigation 1 exit
Figure 3: A sample visualization of Procgen games (left) and Grid Navigation (right).
LATENT STATE
VALUE MODULE
PREDICTED VALUE
Exp
V
h
Vi(i
LATENT STATE
TRANSITION MODULE
PREDICTED LATENT NEXT STATE
LATENT STATE
REWARD MODULE
PREDICTED REWARDS
T
R
h
Ryth, a)
Figure 4: An illustration of the learnalle submodules in Differentiable Tree Search Network
D.1.2 PROCGEN
Procgen is a unique suite consisting of 16 game-like environments, each of which is procedurally generated. This means that they are designed to present slightly altered levels every time they are played. Such design intricacy makes Procgen an ideal choice to test an agent's 's generalization capabilities. It stands in contrast to other commonly used testing suites, like the renowned Atari 2600 games (Mnih et al., 2013) 2016 Badia et al., 2020). The diverse array of environments within Procgen emphasizes the pivotal role of robust policy learning. The environmental diversity in Procgen
20

---

Published as a conference paper at ICLR 2025
softmax
Step2
Step
Step
Rx
Step
Figure 5: An illustration of the Expansion Phase in Differentiable Tree Search Network
Repoet for mooe Expord Bhp
BELIMAN EQUATION
BELLWAN EQUATION
Figure 6: An illustration of the Backup Phase in Differentiable Tree Search Network
underlines the importance of robust policy learning for successful generalization. The open source code for the environments is publicly accessible tt/ht.tp://gg http::/ttttcccco/eeaaaaaaarerooa
21

---

Published as a conference paper at ICLR 2025
Computation Graph Construction
Co A
Backpropagation Gradient Computation
V
Qo
Qa
Figure 7: An illustration of the computation graph construction in Differentiable Tree Search Network
D.2 LEARNING SETUP
D-TSN can serve as drop-in replacement for popular model-free predictors. It can be trained using both online and offline reinforcement learning algorithms. In this work, we employ the offline reinforcement learning (Offfine-RL) framework to focus on the sample complexity and generalization capabilities of D-TSN when compared with the baselines.
Offline-RL. often referred to as batch-RL. is the scenario wherein an agent learns its policy solely from fixed offline dataset of experiences, without further interactions with the environment. Offline-LL framework poses significant challenges, especially when it comes the generalization capabilities of methods, even for seemingly straightforward tasks like navigation. Consider, for instance, that the optimal policy is employed to collect experiences from an environment. It would predominantly select the optimal actions every state. This means that only a fraction of the vast state-action space would be covered in the dataset. As result, the model might not learn about many interactions, such as what happens when collides with a wall, due to the limited training data on such environment interactions. When this world model employed an online search scenario, the search process might unknowingly venture into -fi-tti-iitiiitiion state space. These explorations, stemming from the model's limited generalization capabilities, can lead overly optimistic value predictions, subsequently affecting the v-values computation at the root node. In practical terms, the online search might mistakenly believe it can travel through a wall reach its goal more quickly and would then run into it.
However, the D-TSN design offers solution this problem by jointly optimizing both the world model and the online search. During the training phase, when the online search strays into the out-of-distribution states, might overestimate the value these states. This overestimation would then influence the final output after the search. When such mismatch between the post-search output and the expert action is detected, the gradient descent algorithm adjusts these overestimated values, effectively lowering them to align the final values more closely with the expert action. As direct consequence of this, by the end the training phase, the search process will have effectively learned to ignore these the
states. Even if such expansion does oceur due to e-greddy
22

---

Published as a conference paper at ICLR 2025
exploration during inference, the predicted value for the out-of-distribution state will be small and thus will have negligible impact on the Q-values of the root node due to the max operation during backup.
D.3 TRAINING DATASETS
We employ a separate behavior policy to collect the offline training dataset. Here, the behavior policy can be optimal sub-optimal. An optimal policy generates dataset with lower noise and cleaner training signal, leading to a stable learning process. In contrast, a sub-optimal policy produces noisier dataset, which consequently restricts the quality of the policy that the agent can learn. The selection of the behavior policy depends on domain-specifer requirements and the resources available for data collection. This work explores datasets generated using both optimal (in the Navigation domain) and sub-optimal (in the Procgen domain) behavior policies.
The training dataset, D, consists of trajectories generated using the behavior policy TB. where each trajectory, Tr, is defined as a series of I' tuples, each comprising the state observed, action taken, reward observed, and the corresponding Q-value of the observed state, denoted as:
Ti (St,es Ot Pt.i: Qt,3) t=
The Q-value for state 8t,i can be computed by adding the rewards obtained in the trajectory from timestep & onwards, i.c.
Qt TTT (8t,))
= Ft,i + Te+1. + T142,4 +..,,+..... ...r.
T
Tk
In order to evaluate the sample complexity and generalization capabilities of each method, we collected small number of 1000 trajectories for each test domain for our experiments.
D.3.1 NAVIGATION
For our navigation task, which is relatively small in size, we are able to compute the optimal policy for any given state and configuration. We employ the value iteration algorithm for this purpose, as detailed by |Sutton & Barto (2018). At the beginning of each episode, we formulate a random passage through the central hall. Subsequently, we also randomly determine the starting position of the robot within the hall and its goal position outside of it. We collect a total of 1000 expert trajectories for training. Each of these trajectories incorporates a sequence comprising states, actions, rewards, and Q-values observed throughout the episode.
D.3.2 PROCGEN
When it comes to Procgen, even though there isn't public repository of pre-trained models, there exists an open-source code base for Phasic Policy Gradient (PPG) (Cobbe al. 202T). With this in hand, we could effectively train decent, but sub-optimal, policy for every individual Procgen game starting from the ground up. We rely on the default set of hyperparameters for training a specific policy for each game. Using these policies, we gather sample trajectories for 1000 successful episodes. Just like in the case of Navigation, each of these trajectories represents a sequence comprising states, actions, rewards, and valuus observed throughout the episode.
D.4 LOss FUNCTIONS
As mentioned in the previous section, D-TSN can be trained using various online and offline RL methods. In this work, we are focusing offline-RI RL framework and use Behavior Cloning to train the parameters of D-TSN. Behavior Cloning is a type of supervised learning where the objective is to make the agent's prediction closely approximate the actions taken by the behavior policy in the states collected in the training dataset. To achieve this, we minimize the mean squared error between the predicted and target Q-values. This loss, denoted Lo. is expressed as:
La E(,,, Q.)DD Qo(s,, a.) Q.
(29)
Moreover, during the online search, the transition, reward, and value networks operate on the latent states. Consequently, it is important to ensure that the input to these networks is of consistent scale,
23

---

Published as a conference paper at ICLR 2025
suggested in (Schwarzer al. 2021 Ye al. 2021). To achieve this, apply hyperbolic tangent
(Tanh) normalization on the latent states, thereby adjusting their scale to fall within the range (-1,1).
h tamh(x)
e -
+6 (-1,1
D.4.1 AUXILIARY LOsS FOR
ACTIONS
In the offline-RL setting, there's risk that Q-values for out-of-distribution actions might be overesti- mated (Kumar al. 2020)) as the behavior policy can only cover limited part of the state-action distribution. To address this, incorporate additional CQL (Kumar al., (2020) loss which encourages the agent to adhere to actions observed within the training data distribution. This loss, LD, is defined
CD Ees a:) log Σ cxp (Qo(s), o')) Qo(si, o.))
(30)
D.4.2 AUXILIARY LOSS FOR CONSISTENCY IN THE WORLD MODEL
In order to avoid overburdening the latent states with extraneous information required to reconstruct the original input states, like in Model-based RL methods, we utilize self-suupr consistency loss functions described in (Schwarzer et al. 2021] Ye et al. 202T). These functions aid in maintaining consistency within the transition and reward networks. For example, let us assume state 8; and the subsequent state 8441 resulting from action az. The latent state representations for the environment states 8; and 8;41 can be computed as h; and h4+1 respectively. The latent state encoding h4+1 can be predicted using the transition module, h4+1 = Th, az). We minimize the squared error between the latent representation /4411 and h4+1. to ensure that the transition function To provides consistent predictions for the transitions in the latent space. In accordance with the approach detailed in (Ye et al. (021), we use separate encoding network, referred to as the target encoder, to compute target representations.
LT. Es,,, as, 3441 (fitts - hi+t1) 27
(31)
where
h4t1 = To(hi, as)
h; Es(st)
h441 (8i+1)
The parameters of the target encoder, 0', are updated using an exponential moving average of the parameters of the base encoder, 6, as follows.
8' - a 0 + + a)
Notably, we refrain from adding projection or prediction networks, as done in Schwarzer et al. (2021) and Ye et al. (2021), prior calculating the squared difference did not observe any improvement by doing so.
Further, we also seek to minimize the mean squared error between the predicted reward Re (h, a) and the actual reward observed in the training dataset D.
Lro = Bees on rn (Ro(Ie, a;)
-r:) -rx))
(32)
24

---

Published as a conference paper at ICLR 2025
D.4.3 FINAL Loss FUNCTION
The final loss function train D-TSN is a combination of Behavior Cloning loss and all the auxiliary losses defined above. It is given by:
L = E, Ailo + 22CD) + A3CTo + AACRo
(33)
where A1, A2, λ3 and d serve as weighting hyperparameters.
D.5 BASELINES
To evaluate the efficacy of Tree Search Network, benchmark it against the following prominent baselines:
Model-free Q-network: This allows to assess the significance of integrating the inductive biases into the neural network architecture.
Model-based Search: In this baseline, we assess the benefits derived from the joint opti- mization of the world model and the search algorithm by training the world models and value module independently of each other and utilizing them for online search.
TreeQN: This comparison helps in highlighting the advantages of using a more advanced search algorithm that can execute deeper search while maintaining similar computational constraints.
D.6 IMPLEMENTATION DETAILS
In an effort to assess the distinctive elements of each method's s design, we ensure uniformity in the number of parameters across all agents. This is achieved by integrating the submodules from D-TSN into the network architecture each baseline. However, while the number of parameters are consistent, the way in which these submodules are utilized to construct the computation graph varies among the baselines. provide their implementation details below:
D.6.1 DIFFERENTIABLE TREE SEARCH NETWORK
D-TSN utilizes its submodules in alignment with the best-first search algorithm presented in Sec- For our empirical evaluations, we set the maximum limit for search iterations at 10. Throughout the training process, the computation graph, formulated via online search, optimized to accurately predict the Q-values. This optimization serves a dual purpose: it not only refines the Q-value predictions but also facilitates robust learning for the submodules when they are employed in context of online search. The loss function utilized for training is:
tion 3.2
CD-TSN = E. Aifo + >₂DD + dghTo + 444
(34)
We compute the gradient for this loss as described in Equation
D.6.2 MODEL-FREE Q-NETWORK
In this baseline, the submodules are utilized to perform one-step look-ahead search. The input state at the root node undergoes an expansion using the world model, and Q-values are computed using the Bellman equation represented Q(s, a) Rew(h, a) Val(h'), where h Enc(s) and hí = Tr(h, a). Intriguingly, this structure does encapsulate basic inductive bias via the I-step look-ahead search. However, in keeping with its model-free characteristic, auxiliary losses aren't 'ttt employed for training the transition and reward model. The loss function for this model is:
LoNet = Xılo + dgL
D.6.3 MODEL-BASED SEARCH
For this approach, we employ the best-first algorithm showcased in Algorithm [ However, there's difference: the world model and the value module are trained independently, each focusing on their specific objectives. As outlined in Ye al. (202I). we incorporate
supervised consistency losses defined in Equations (31 and (32) as they improve the online search, even in cases where the world model not jointly trained with the online search. The Q-values used for training are computed directly using the value module without performing online search during training. The loss function used for this approach is:
Lsearch -iia + DoLo >gLTo + dsLR..
25

---

Published as a conference paper at ICLR 2025
D.6.4 TREEQN
In this baseline, the starting step is encoding the input state its latent counterpart with the Encoder module. Following this, a full-tree expansion, based on predefined depth d, is performed using both Transition and Reward modules. The values at the leaf nodes are then backpropagated to the root node via the Bellman equation, discussed in the Backup phase in Section 3.2 The root node Q-values serve as the final output, that is utilized for training. Given the exponential growth of TreeQN's computation graph with an increase in depth d, we choose a depth of 2 for both Procgen and navigation domains in our analysis, as used in TreeQN's original code base (Farquhar et al. 2018). Notably, greater depths, such as 3 or more, are infeasible since the resulting computation graph exceeds the memory capacity (roughly 11GB) of a typical consumer-grade GPU. The associated loss function, as adapted from the original paper, is:
LTeeQQ X + A2CD + AgCR,
is important to note that every method is trained with same datasets using their respective loss functions. We fine-tune the hyperparameters, A1, 22, λ3 and Y1, using grid search on a log scale.
D.7 ISSUE WITH BASELINE NORMALIZED SCORE
Within the context of Atari games, the Baseline Normalized Score (BNS) is frequently utilized to evaluate the performance of agents. When human players are used as the baseline, is often termed Human Normalized Score. The primary allure of BNS lies in its capacity to offer a relative assessment of an agent's 's performance, comparing it against a standard benchmark this could be human players or even another agent.
One of the primary benefits of the BNS is its ability to provide a consistent metric across different games, addressing the difference in scale inherent in raw scores. By enabling the calculation of the average BNS across multiple games, we gain insight into the overall efficacy of an agent. This not only facilitates direct performance comparisons between diverse agents and methodologies but also paints a picture of how the agent's 's abilities stack up against human standards.
To derive the BNS, start by logging the agent's then normalized against a baseline score, derived from baseline agent's performance on the same game. By dividing the agent's score by the baseline's 's score (and sometimes subtracting the score of random agent), we get a relative metric. Mathematically, this can be represented as:
score in Atari game. This raw score is
Sp-SR SRR
BNS(m)
(35)
SB SS
Here, Sz,, SB and Sn denote the raw scores of the agent, the baseline policy, and random policy, respectively. Interpretation-wise BNS of 1 indicates parity with the baseline. Values exceeding signify that the method outperforms, while those below I indicate that the method underperforms relative to the baseline.
Nevertheless, the BNS has its frailties. It inherently presumes the baseline policy will always surpass the performance of the random policy. But there can be instances, contingent on the environment the specific baseline policy, where this isn't the case. In scenarios where the bascline policy underperforms the random policy, the BNS results in a negative denominator. This poses a predicament: even if our agent's 's policy performs better than the random policy, the BNS unfairly penalizes it. In our experiments with Procgen, we observed that for 2 out of the 16 games, namely heist and maze, the baseline policy underperformed compared to the random policy. Given these pitfalls, our evaluations pivot towards more robust metric: the Z-score. The Z-score, often termed as the "standard score". provides statistical measurement that describes value's 's relationship the mean of a group of values. It is measured in terms of standard deviations from the mean. If a Z-score O, indicates that the data point's 's score identical to the mean score. Z-scores may be positive or negative, with a positive value indicating the score is above the mean and a negative score indicating is below the mean.
26


Reviews:
[{"review_id": "wV7KGicMxH", "reviewer": "Reviewer_tL81", "summary": "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.", "questions": "1. How would you extend D-TSNs to use for solving stochastic decision making problems? What modification would be required in the current work to accommodate stochastic transitions and rewards?\n\n2. Could you provide more details on how computationally expensive D-TSN is and how does the method scale w.r.t. the action space size and the search tree depth?", "limitations": "", "rating": 10, "confidence": 4, "soundness": 4, "presentation": 4, "contribution": 3, "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. \n\n2. 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. \n\n3. 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. \n\n4. 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. \n\n5. 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.\n\n2. 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.\n\n3. Scalability: scalability is not thoroughly analyzed for longe-horizon tasks or higher dimensional state space.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "LGOLIOSJtR", "reviewer": "Reviewer_7dAQ", "summary": "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.", "questions": "- Why did the authors choose to use REINFORCE? Did the authors experiment with using methods like Gumbel-Softmax for sampling?\n- How are the compute budgets between model-based search, model-free Q network and D-TSN equalized for a fair comparison?\n- In Line 22, the authors say they “introduce an effective variance reduction technique”. Isn’t this just Guez et al. (2018)? The way this is written in some places suggests that this is novel.\n- Line 45,46 seems contradictory with lines 48,49. 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. I understand the CQL term attempts to address this issue.\n- What are the scores of the PPG sub-optimal policy?\n- Table 3 reports Mean Scores and Mean Z-Scores, but no standard deviations or error bars?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 4, "presentation": 2, "contribution": 3, "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.\n- Strong empirical evidence that the proposed method improves on TreeQN, and having the modules trained separately.\n- 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.\n- The proposed method is mostly applicable to discrete action spaces with deterministic environments. Improving on this remains a future empirical question.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "RBacLky96F", "reviewer": "Reviewer_KW9a", "summary": "This paper introduces a differentiable tree search and planning network to alleviate the suboptimal search results that arise when using only action and observation sequences to learn to plan in many modern data-driven planning problems. The tree search basically synergizes submodules in the form of an encoder, value, reward and transition functions from a network by inducing the bias of a best first search into the neural architecture.\n\nArguing that for a slight perturbation in network parameters, the implementation of the loss function in equation (2) could generate a tree structure that causes the loss function to become noisy, the authors equated this to a lack of continuity in the loss function space. I think they mistook stochasticity in gradient propagation with discontinuity. The whole premise of the contribution of the paper is based on this assumption that is barely proven to be true or false before the authors dived into a host of mathematical analysis that resulted in the loss function on Line 272 (please number all your equations to allow for easy referencing in the future). As a result of this key oversight, it is not clear to this reviewer that the whole contribution of the paper is warranted. \n\nAs a result, I am assigning a score of fair to this contribution until I see a well-laid out argument for why this new invention is necessary.", "questions": "+ How is the full expansion mechanism in this work not a shallow tree (in comparison to the TreeQN) that you mentioned in the RELATED WORKs section since you mentioned that you have a fixed depth in your search protocol?\n\n+ When the encoder, value, reward, and transition functions are unavailable, how do they get jointly optimized with the search algorithms?\n\n+ In paragraph III, the justification for why end-to-end learning improves the reliability of world-models used in preconditioning the search process is not given. Please address this.\n\n+ Lines 124-126: \"Dividing the network into these submodules reduces\ntotal 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.\" Where is the proof for this?\n\n+ Is there no way the illustration in Figure 2 can be moved to the main text from the Appendix? Probably as a subfigure on page one or so would be really nice! Also, I think the algorithm should justifiably be compressed (e.g. as a pseudocode or flow chart) somewhere in the main text. An elaborate version can be embedded in the Appendix if space is an issue.\n\n\n+ What metric are you using to report the measured quantities in Table 1? Can you include it in the heading/sub-headings?\n\n\n+ 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? \n\n+ Line 915, I think you meant Lemma B.3, not Theorem B.3, no?", "limitations": "", "rating": 8, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 4, "strengths": "+ This paper introduces a differentiable treee seach and planning network to alleviate the suboptimal search results that arise when using only action and observation sequences to learn to plan in many modern data-driven plannning problems. The tree search basically synergizes submodules in the form of an encoder, value, reward and transition functions from a network by inducing the bias of a best first search into the neural architecture.\n\n+ I love the motivation stated for constructing an search tree expansion policy that is stochastic in nature. But the justification for why it ensures the continuity of the loss function when the search tree is not fixed is missing. I am referring to lines 55-56.\n\n+ I love the conceptualization and the synergy of REINFORCE, mathematical mechanisms to reduce variance in the REINFORCE Loss owing to possibly biased estimates, the continuity proof (though I have questions hanging over the proof of Lemma B.3 to be fully satisfied with this poposition) .\n\n+ I love that the conclusion section meticulously summarizes the problem, contributions, and shortcomings. Kudos to the authors.", "weaknesses": "While I do love the mathematical contributions of the paper, I think there are essentials that are missing in the logic, organization, and flow of arguments that require a thorough review before this paper makes it into an acceptance. A principal one is the following (mentioned in the summary box but repeated here. The authors would do well to alleviate my concerns): Arguing that for a slight perturbation in network parameters, the implementation of the loss function in equation (2) could generate a tree structure that causes the loss function to become noisy, the authors equated this to a lack of continuity in the loss function space. I think they mistook stochasticity in gradient propagation with discontinuity. The whole premise of the contribution of the paper is based on this assumption that is barely proven to be true or false before the authors dived into a host of mathematical analysis that resulted in the loss function on Line 272 (please number all your equations to allow for easy referencing in the future). \n\n+ The claim in the last paragraph of Theorem 3.1 that slight changes in the network parameters could cause discontinuity in the tree structure seems anecdotal and not backed up by a solid proof. I would love to see a concrete reasoning (analytical proof or abundant empirical  proofs) behind this claim that warrants section 3.5 and Appendix C.\n\n+ Grammatical errors fly out of the page hither and yon throughout the paper; the uthors would do well to carefully organize their arguments, present their logic convincingly throughout the paper, and punctuate and label every equation appropriately!\n\n+ The logic in the paper could use a more thoughful presentation. Here is an example critique:\n     - In the \"introduction\", it is stated in the first paragraph that constructing a search tree from expert data is infeasible due to lack of practicality or scarcity. The authors make an assumption that a search tree is a principal prerequisite for information retrieval (IR) without any justification as to why it may be better than alternative IR methods. Then in the second paragraph, they mentioned how search and planning could be better executed in the presence of a simulator or world mode. While I find this premise alluring, 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}, {"review_id": "lMFexNa8op", "reviewer": "Reviewer_iWz4", "summary": "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.", "questions": "Do differences in the convergence rates of submodules exist? If present, could these differences impact the overall network's performance?", "limitations": "", "rating": 6, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "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.\n\nii) 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. \n\niii) 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.\n\n\nii) 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.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "", "claims_and_evidence": "", "methods_and_evaluation_criteria": "", "theoretical_claims": "", "experimental_designs_or_analyses": "", "relation_to_broader_scientific_literature": "", "essential_references_not_discussed": "", "other_strengths_and_weaknesses": "", "other_comments_or_suggestions": "", "questions_for_authors": "", "summary_of_paper": "", "summary_of_recommendation": "", "technical_quality": "", "clarity_of_presentation": "", "potential_impact": "", "robotics_focus": "", "quality_of_limitations": "", "questions_for_rebuttal": "", "reasons_to_accept": "", "reasons_to_reject": "", "questions_to_authors": "", "extra_scores": {}}]

Author rebuttal:

Decision: Accept (Oral)