STAMP: Differentiable Task and Motion Planning via Stein Variational Gradient Descent

TL;DR

This work proposes a novel approach to TAMP called Stein Task and Motion Planning (STAMP) that relaxes the hybrid optimization problem into a continuous domain and solves the optimization problem using a gradient-based variational inference algorithm called Stein Variational Gradient Descent.

Abstract

Planning for sequential robotics tasks often requires integrated symbolic and geometric reasoning. TAMP algorithms typically solve these problems by performing a tree search over high-level task sequences while checking for kinematic and dynamic feasibility. This can be inefficient because, typically, candidate task plans resulting from the tree search ignore geometric information. This often leads to motion planning failures that require expensive backtracking steps to find alternative task plans. We propose a novel approach to TAMP called Stein Task and Motion Planning (STAMP) that relaxes the hybrid optimization problem into a continuous domain. This allows us to leverage gradients from differentiable physics simulation to fully optimize discrete and continuous plan parameters for TAMP. In particular, we solve the optimization problem using a gradient-based variational inference algorithm called Stein Variational Gradient Descent. This allows us to find a distribution of solutions within a single optimization run. Furthermore, we use an off-the-shelf differentiable physics simulator that is parallelized on the GPU to run parallelized inference over diverse plan parameters. We demonstrate our method on a variety of problems and show that it can find multiple diverse plans in a single optimization run while also being significantly faster than existing approaches.

Figures (12)

• Figure 1: Demonstrations of the top-3 block-pushing plans found by STAMP on a Franka manipulator. Task plans from left to right: push the East side; push the East side, then the North side; push the North side. STAMP found all three solutions in one run.
• Figure 2: Overview of STAMP algorithm pipeline. The particles $\theta$ represent the task and motion plan $\theta=[a_{1:K}, u_{0:KT-1}]$, or the task plan and a parametrization of the motion plan $\theta=[a_{1:K}, g_{1:K}]$ if Dynamic Movement Primitives dependent on goals $g_{1:K}$ are used (see section \ref{['sec:dmp_for_stamp']}). The resulting controls are passed through a differentiable physics simulator, and particles are updated in parallel using two phases of optimization: an inference phase of SVGD that balances loss minimization with diversification of particles followed by a finetuning phase of SGD so that particles can reach the local optima.
• Figure 3: Fitting a Gaussian mixture with discrete-and-continuous SVGD. Left: random initialization of particles. Right: particles after 500 updates. Note that same-color particles have been assigned to the same mode.
• Figure 4: Sample solutions STAMP found for the billiards problem. Subcaptions indicate which walls the cue ball hits before hitting the target ball. The $1^\text{st}$/$2^\text{nd}$/$3^\text{rd}$/$4^\text{th}$ digit correspond to hitting the top/bottom/left/right wall. A 1 means the wall is hit during the shot; a 0 means that it is not. The caption also indicates in which pocket the target ball is shot.
• Figure 5: Evolution of mean cost over all particles averaged across 5 runs.