A Unified Framework for Planning with Learned Neural Network Transition Models

Automated planning with neural network transition models is a two stage approach to solving planning problems with unknown transition models. The first stage of the approach learns the unknown transition model from data as a neural network model, and the second stage of the approach compiles the learned model to either a Mixed-Integer Linear Programming (MILP) model or a Recurrent Neural Network (RNN) model, and optimize it using an off-the-shelf solver. The previous studies have shown that both models have their advantages and disadvantages. Namely, the MILP model can be solved optimally using a branch-and-bound algorithm but has been experimentally shown not to scale well for neural networks with multiple hidden layers. In contrast, the RNN model can be solved effectively using a gradient descent algorithm but can only work under very restrictive assumptions. In this paper, we focus on improving the effectiveness of solving the second stage of the approach by introducing (i) a novel Lagrangian RNN architecture that can model the previously ignored components of the planning problem as Lagrangian functions, and (ii) a novel framework that unifies the MILP and the Lagrangian RNN models such that the weakness of one model is complemented by the strength of the other. Experimentally, we show that our unifying framework significantly outperforms the standalone MILP model by solving 80% more problem instances, and showcase the ability of our unifying framework to find high quality solutions to challenging automated planning problems with unknown transition models.