Differentiable Tree Search in Latent State Space
Dixant Mittal Wee Sun Lee
In decision-making problems with limited training data, policy functions
approximated using deep neural networks often exhibit suboptimal performance.
An alternative approach involves learning a world model from the limited data
and determining actions through online search. However, the performance is
adversely affected by compounding errors arising from inaccuracies in the
learnt world model. While methods like TreeQN have attempted to address these
inaccuracies by incorporating algorithmic structural biases into their
architectures, the biases they introduce are often weak and insufficient for
complex decision-making tasks. In this work, we introduce Differentiable Tree
Search (DTS), a novel neural network architecture that significantly
strengthens the inductive bias by embedding the algorithmic structure of a
best-first online search algorithm. DTS employs a learnt world model to conduct
a fully differentiable online search in latent state space. The world model is
jointly optimised with the search algorithm, enabling the learning of a robust
world model and mitigating the effect of model inaccuracies. We address
potential Q-function discontinuities arising from naive incorporation of
best-first search by adopting a stochastic tree expansion policy, formulating
search tree expansion as a decision-making task, and introducing an effective
variance reduction technique for the gradient computation. We evaluate DTS in
an offline-RL setting with a limited training data scenario on Procgen games
and grid navigation task, and demonstrate that DTS outperforms popular
model-free and model-based baselines.