Keyi Shen^1*, Jiangwei Yu^1*, Jose Barreiros², Huan Zhang¹, Yunzhu Li³ (* indicates equal contributions)

¹University of Illinois Urbana-Champaign ²Toyota Research Institute ³Columbia University

• [Dec 2024] Initial release for reproducibility. More accessible codebase will be released with α,β-CROWN in 2025.

This repository contains the implementation of BaB-ND, a GPU-accelerated branch-and-bound (BaB) framework for motion planning in manipulation tasks that require trajectory optimization over neural dynamics models (ND).

To solve the complex planning problem involving neural dynamics, our approach BaB-ND, divides the action space into smaller sub-domains through a novel branching heuristic (branching), estimates objective bounds using a modified bound propagation procedure, inspired by the state-of-the-art neural network verifier α,β-CROWN (alpha-beta-CROWN), to prune sub-domains that cannot yield better solutions (bounding), and focuses searches on the most promising sub-domains (searching). The branching process guides searching effectively, while the bounding process strategically reduces the search space.

We demonstrate the effectiveness, applicability, and scalability of our framework across a range of complex planning problems, including contact-rich manipulation tasks, the handling of deformable objects, and object piles, using diverse model architectures such as multilayer perceptrons and graph neural networks.

Please refer to our paper for more details. This repository provides the code to genreate some qualitative exmaples and reproduce some quantitative results in the paper. The main algorithm is implemented based on the neural network verifier α,β-CROWN in the Verifier_Development directory and will be officially published in the future release of α,β-CROWN.

1. We work on the Ubuntu 22.04.5 LTS operating system and use Miniconda to manage our environment and required packages. Please install Miniconda if needed: https://docs.conda.io/en/latest/miniconda.html.
2. Create the environment by running the following command. It will create a new environment named bab-nd with Python 3.9.17, Torch 2.1.0, and CUDA 12.1. We tested the environment and code on a NVIDIA GeForce RTX 4090.

   conda env create -f environment.yaml
   

3. Activate the environment.

   conda activate bab-nd
   

4. Install the associated package for the Neural Network Verifier α,β-CROWN.

   cd Verifier_Development
   pip install -e .
   cd ..
   

1. Run the following command to perform a quick test on the synthetic example:

   python tasks/func_test/func_test.py
   

   This will execute 4 methods (GD, MPPI, CEM, and ours) on the synthetic example with an input dimension of 100. Only one random seed is used in the test.

   The test generally takes less than 1 minute to complete on an NVIDIA GeForce RTX 4090, and the results will be saved in output/func_test/quick_test.txt. The visualization of the function in 1D will be stored in output/func_test/function_plot_1d.pdf.

   A sample output can be found in output_sample/func_test/quick_test.txt, and the corresponding function plot is in output_sample/func_test/function_plot_1d.pdf.

2. Run the following command to perform the complete test on the synthetic example:

   python tasks/func_test/func_test.py -complete
   

   This will execute 4 methods (GD, MPPI, CEM, and ours) on the synthetic example with input dimensions ranging from 5 to 100, in intervals of 5. Only one random seed is used in the test.

   The test generally takes less than 10 minutes to complete on an NVIDIA GeForce RTX 4090, and the results will be saved in output/func_test/complete_test.pdf.

   A sample output can be found in output_sample/func_test/complete_test.pdf or below.

Run the following command to try this task:

bash scripts/pushing_w_obs_qualitative.sh

This will run 4 methods (GD, MPPI and CEM, and ours) on 3 cases with different configurations (initial state, target state, and obstacle setting).

• Note: To reduce the requirements for computational resources and runtime, we optimize the final step cost considering the obstacles and use a smaller number of samples and iterations for all methods.

The test generally takes about 4 minutes to complete on an NVIDIA GeForce RTX 4090. The results will be stored in output/pushing_T/qualitative/. There will be 3 sub-directories for the 3 cases. Each sub-directory contains the results of the 3 methods. Inside each sub-directory:

• summary.txt provides the objective of open-loop planning, the final cost after closed-loop execution, and the runtime of open-loop planning for the 3 methods.
• long_horizon_planning_0.gif shows the trajectories of long-horizon open-loop planning from the 3 methods.
• execution_traj_0.gif shows the trajectories of closed-loop execution following the open-loop trajectories.
• (optional) experiment_results.json contains detailed information about the experiment.
• Note: In the results, GD is MPPI displayed as "GD_BF", is displayed as "MPPI_BF", CEM is displayed as "DecentCEM" and our method is displayed as "CROWN".

A sample output can be found in output_sample/pushing_T/qualitative/. Below are the execution trajectories of our BaB-ND in the 3 cases.

Run the following command to reproduce the quantitative results of this task:

bash scripts/pushing_w_obs_quantitative.sh

The test generally takes about 80 minutes to complete on an NVIDIA GeForce RTX 4090. The results will be stored in output/pushing_T/quantitative/. planning_stat.pdf will be the same as the subplot Pushing w/ Obstacles of Figure 6 (a) Open-loop Performance in the paper.

A sample output can be found in output_sample/pushing_T/quantitative/ or below.

Run the following command to try this task:

bash scripts/object_sorting_qualitative.sh

This will run 4 methods (GD, MPPI and CEM, and ours) on 3 cases with different configurations (initial state and target state).

• Note: To reduce the requirements for computational resources and runtime, we optimize the final step cost with a shorter planning horizon and use a smaller number of samples for all methods.

The test generally takes about 4 minutes to complete on an NVIDIA GeForce RTX 4090. The results will be stored in output/pushing_T/qualitative/. There will be 3 sub-directories for the 3 cases. Each sub-directory contains the results of the 3 methods. Inside each sub-directory:

• summary.txt provides the final cost after closed-loop execution, and the total runtime of planning during closed-loop execution for the 3 methods.
• execution_traj_0.gif shows the trajectories of closed-loop execution.
• (optional) experiment_results.json contains detailed information about the experiment.
• Note: In the results, GD is MPPI displayed as "GD_BF", MPPI is displayed as "MPPI_BF", CEM is displayed as "DecentCEM" and our method is displayed as "CROWN".

A sample output can be found in output_sample/pushing_T/qualitative/. Below are the execution trajectories of our BaB-ND in the 3 cases.

Run the following command to reproduce the quantitative results of the pushing with obstacles task:

bash scripts/object_sorting_quantitative.sh

The test generally takes about 100 minutes to complete on an NVIDIA GeForce RTX 4090. The results will be stored in output/obj_pile/quantitative/. planning_stat.pdf will be the same as the subplot Object Sorting of Figure 6 (a) Open-loop Performance in the paper.

A sample output can be found in output_sample/obj_pile/quantitative/ or below.

If you find our project helpful, please cite our paper:

@misc{shen2025babndlonghorizonmotionplanning,
      title={BaB-ND: Long-Horizon Motion Planning with Branch-and-Bound and Neural Dynamics}, 
      author={Keyi Shen and Jiangwei Yu and Jose Barreiros and Huan Zhang and Yunzhu Li},
      year={2025},
      eprint={2412.09584},
      archivePrefix={arXiv},
      primaryClass={cs.RO},
      url={https://arxiv.org/abs/2412.09584}, 
}