Scaling up and Stabilizing Differentiable
Planning with Implicit Differentiation

We derive implicit differentiation for VIN and variants, where each layer $f$ is a step as in value iteration. Since the derivation does not rely on the concrete instantiation of $f$, we can freely replace $f$ from Conv2D with other types of layers. We will introduce these variants in the next subsection.

Implicit differentiation. A fixed-point problem can be solved iteratively by fixed-point iteration and other algorithms. However, as pointed out in (Bai et al., 2019), naively differentiating through the solver would require intensive memory usage, since it needs to store every intermediate iterate $V^{(k)}$ to enable automatic differentiation, as in (Tamar et al., 2016; Lee et al., 2018). As also used recently in (Bai et al., 2019; Nikishin et al., 2021; Gehring et al., 2021), another solution is to instead differentiate directly through the fixed point $z^{*}$ using the implicit function theorem and implicit differentiation. Then, we can skip all of this by decoupling forward (fixed-point iteration as the solver) and backward pass (differentiating through the solver).

We start with the fixed point equation ${\bm{v}}^{\star}={\bm{f}}({\bm{v}}^{\star},{\bm{r}},{\bm{\theta}})$ from the Bellman optimality equation. Below we use ${\bm{x}}$ to stand for either input ${\bm{r}}$ or ${\bm{\theta}}$. The implicit function theorem tells us that, under some mild conditions of the derivatives (${\bm{f}}$ is continuously differentiable with non-singular Jacobian $\nicefrac{{\partial{\bm{f}}({\bm{v}},{\bm{x}})}}{{\partial{\bm{x}}}}$), ${\bm{v}}^{\star}({\bm{x}})$ is a differentiable function of ${\bm{x}}$ locally: $0=f({\bm{v}}^{\star}({\bm{x}}),{\bm{x}})$.

For fixed point equation, we can assume the fixed point solution ${\bm{v}}^{\star}$ is obtained, thus this can be used to compute the partial derivative w.r.t. to any quantity (input, parameter, etc) $\nicefrac{{\partial{\bm{v}}^{\star}(\cdot)}}{{\partial(\cdot)}}$. It avoids backpropagating gradients through the forward fixed-point iteration, which is computationally inefficient and requires considerable memory to store intermediate iteration variables. Additionally, it also allows to even use of non-differentiable operations in the forward pass.

Differentiating both sides of the equation $\bm{v}^{\star}=f(\bm{v}^{\star},\bm{x})$ and applying the chain rule:

where we use $(\cdot)$ to denote an arbitrary variable. Rearranging terms:

Solving backward pass. To integrate into a deep learning framework for automatic differentiation, two quantities are needed: VJP (vector-Jacobian product) and JVP (Jacobian-vector product) (Gilbert, 1992). Nevertheless, the computation of the inverse term $(I-\nicefrac{{\partial f(\bm{v}^{\star},\bm{x})}}{{\partial\bm{v}^{\star}}})^{-1}$ can be a major bottleneck due to its dimension ($O(d^{2})$ or $O(m^{4})$, where $d=m^{2}$ is the matrix width and $m$ is the map size) (Bai et al., 2019; 2021). Additionally, when applied to VINs, we concatenate a policy layer that maps the final equilibrium $\bm{v}^{\star}\in\mathbb{R}^{m\times m}$ to action logits $\mathbb{R}^{m\times m}$ and compute the cross-entropy loss $\ell(\cdot)$ (Tamar et al., 2016). Thus, the derivative of loss is:

Defining $\bm{w}$ as follows forms a linear fixed-point system (Bai et al., 2019):

This backward pass fixed-point equation can also be solved by a generic fixed-point solver or root finder. Then, we can substitute the solution $\bm{w}$ back: $\nicefrac{{\partial\ell}}{{\partial(\cdot)}}=\bm{w}^{\top}\nicefrac{{\partial f(\bm{v}^{\star},\bm{x})}}{{\partial(\cdot)}}$. The computation is then purely based on VJP and JVP. In summary, an IDP needs to solve both the (nonlinear) forward fixed-point system and the (linear) backward fixed-point system, as in DEQ.

We can derive variants of VIN using implicit differentiation by abstracting out the implementation of value iteration layer. In this section, we propose a generic implicit planning pipeline to extend our approach to Gated Path Planning Networks (GPPN) (Lee et al., 2018) and Symmetric VIN (SymVIN) (Zhao et al., 2022). Spatial Planning Transformers (SPT) (Chaplot et al., 2021) also fits into the pipeline, but it performs less well, as discussed in Section C.

Figure 3 shows the general pipeline, where the network layer $\bm{f}_{\theta}$ can be replaced by any single layer that is capable of iterating values. The pipeline follows VIN and GPPN, where for 2D path planning a map $\mathbb{Z}^{2}\to\{0,1\}$ is provided, and the planners’ output actions (their logits) for each position $\mathbb{Z}^{2}\to\mathbb{R}^{\mathcal{|A|}}$. All these tasks can be represented as taking a form of map “signal” over grid $\mathbb{Z}^{2}$, as previously been done (Zhao et al., 2022; Chaplot et al., 2021).

Planner instantiations. We now introduce the instantiations of implicit planners one by one. We focus on the value iteration part (omit map input and action output), and all planners follow the form $\bar{V}^{(k+1)}=f_{\bm{\theta}}(\bar{V}^{(k)},\bar{R})$ (bars omitted later). There are two design principles: (1) input inject ($\bar{R}$ must be input, as input ${\bm{x}}$) and (2) weight-tied (${\bm{\theta}}$ is shared across layers $f$), as also used in DEQ (Bai et al., 2019). Specifically, the purpose of input inject is that fixed-point solution $\bar{V}^{\star}$ does not depend on initialization $\bar{V}^{(0)}$, so we must pass information of the map through input inject (by $\bar{R}$).

(i) ID-VIN uses regular 2D translation-equivariant convolution layer, where $f(V,R)=\max_{a}R^{a}+\texttt{Conv2D}(V;\theta)$. (ii) ID-SymVIN aims to integrate symmetry into planning and uses equivariant $E(2)$-steerable CNN (Weiler & Cesa, 2021). It has similar form to VIN and just replaces Conv2D with SteerableConv, thus the form is $f(V,R)=\max_{a}R^{a}+\texttt{SteerableConv}(V;\theta)$.

(iii) ID-ConvGPPN is based on our modified version of GPPN (Kong, 2022; Zhao et al., 2022), where we (1) use GRU since it has a single input with form $z^{\prime}=\text{GRU}(z,x)$ and is easier to integrate into our current form, (2) replace all fully connected layers with convolution layers, and (3) inject $R$ to every step. The result is that every layer is a ConvGRU, instead of LSTM in GPPN: $f(V,R)=\texttt{ConvGRU}(V,R;\theta)$. Note that the GPPN variants do not have $\max$ in each iteration anymore and directly take reward $R$ to the recurrent cell (Lee et al., 2018).

Mapper Layer. We can handle tasks with more challenging input, such as visual navigation and workspace manipulation (Lee et al., 2018; Chaplot et al., 2021; Zhao et al., 2022), by learning an additional mapping network (mapper) to first map the input to a 2D map. Further details about environments and mapper implementation are deferred to Section 5.1 and Section C.

Optimization. We build upon the deep equilibrium model (DEQ) and relevant work (Bai et al., 2019; 2020; 2021), which includes several effective techniques. By representing the implementation of value iteration as fixed-point solving, we have the flexibility to use any fixed-point or root solver for $f({\bm{v}},{\bm{\theta}})={\bm{v}}$ or ${\bm{v}}-f({\bm{v}},{\bm{\theta}})=0$. A straightforward solver is forward iteration, as used in VIN’s feedforward network (Tamar et al., 2016). However, recent work has employed Anderson’s or Broyden’s method (Bai et al., 2019; 2020). In our experiments, we compare the use of forward iteration and the Anderson solver in both forward and backward passes. Notably, the SymVIN architecture requires the entire forward pass to be equivariant, so extra attention is necessary when designing the forward solver. Further details and results can be found in Sections C and E.

Underlying computational similarity. The gradient computation is done by automatic differentiation (Gilbert, 1992). For algorithmic differentiation, the gradients are computed through direct backpropagation and the implementation is also based on efficiently computing vector-Jacobian product (VJP) and Jacobian-vector product (JVP) (Bai et al., 2019). Christianson (1994) studied automatic differentiation for implicit differentiation of fixed-point system. The only difference is the number of operations required and that implicit differentiation is based on the Jacobian at the equilibrium. We derive the connection in Section B.1.

Comparison: Implicit vs. algorithmic differentiation. In Figure 4, we compare the performance of three IEPs with different numbers of iterations in the forward pass. Unlike the ADPs shown in Figure 2, our IDPs converge stably with larger forward-pass iterations (horizontal axis), while ADPs sometimes diverge on large iterations. Note that we should not directly compare IDPs and ADPs with the same number of forward iterations since they refer to different things: IDPs compute an equilibrium and use it for backward pass (more forward iterations would be better), while for ADPs # iterations = # layers (more iterations/layers leads to instability). Further analyses in Section 5.2.

Tradeoff: The quality of equilibrium. Theoretically, IDPs have a constant cost of the backward pass with respect to the forward planning horizon. However, the backward pass requires the Jacobian of the final equilibrium $\bm{v}^{\star}\approx\bm{v}_{K}$ from the forward pass. If the equilibrium $\bm{v}^{\star}$ is not solved reasonably well, the backward pass in Eq. 6 would iterate based on an inaccurate Jacobian $\nicefrac{{\partial f(\bm{v}^{\star},x)}}{{\partial\bm{v}^{\star}}}$, which would cause poor performance. In contrast, because ADPs compute exact gradients by backpropagation through layers, they do not suffer from this issue. Additionally, fewer iterations (layers) would also alleviate their convergence issues.

Empirically, with fewer forward iterations (e.g., 10), we observe a performance drop from IDPs. Although ADPs also perform worse with fewer layers, they have a smaller drop compared to IDPs. However, when scaling up to more forward iterations $K_{\text{fwd}}$ (e.g. $\geq 30$ iterations, which are necessary for maps $\geq 27\times 27$) IDPs may solve the equilibrium well enough and are more favorable because of their efficient backward pass. Moreover, we find that using around $K_{\text{bwd}}\approx 15$ iterations for backward pass works well enough consistently across different map sizes ($15\times 15$ through $49\times 49$). Since both algorithmic and implicit differentiation use a similar amount of vector-matrix products, using more than $K_{\text{layer}}\geq 15\approx K_{\text{bwd}}$ would consume more resources and favor IDPs.

Environments and datasets. We run implicit and algorithmic differentiable planners on four types of tasks: (1) 2D navigation, (2) visual navigation, (3) 2 degrees of freedom (2-DOF) configuration space manipulation, and (4) 2-DOF workspace manipulation. These tasks require planning on either given (2D navigation and 2-DOF configuration-space manipulation) or learned maps (visual navigation and 2-DOF workspace manipulation), where the maps are 2D regular grid as in prior work (Tamar et al., 2016; Lee et al., 2018; Chaplot et al., 2021). To learn maps, a planner needs to jointly learn a mapper that converts egocentric panoramic images (visual navigation) or workspace states (workspace manipulation) into a 2D grid. We follow the setup in (Lee et al., 2018; Chaplot et al., 2021) and further discuss in Section D. In both cases, we randomly generate training, validation and test data of $10K/2K/2K$ maps for all map sizes. For all maps, the action space is to move in 4 $\circlearrowleft$ directions: $\mathcal{A}=\texttt{(north,west,south,east)}$.

Training and evaluation. We report success rate and training curves over $5$ seeds. The training process (on given maps) follows (Tamar et al., 2016; Lee et al., 2018; Zhao et al., 2022), where we train $60$ epochs with batch size $32$, and use kernel size $F=3$ by default. We use RTX 2080 Ti cards with 11GB memory for training, thus we use 11GB as the memory limit for all models.

In the previous sections with Figure 2 and 4, we have presented quantitive analysis of IDPs and ADPs in terms of convergence with more iterations and scalability on larger tasks. Here, we provide the detailed setup of the experiment and put the attention more on the qualitative side.

Setup. We train all models on 2D maze navigation tasks with map sizes $15\times 15$, $27\times 27$, and $49\times 49$. We use $K_{\text{layer}}=30,50,80$ iterations for ADPs, which is effectively the number of layers in their networks. Correspondingly, for IDPs, we choose to use forward iteration solver for forward pass and Anderson solver for backward pass. We fix the number of iterations of backward solver as $K_{\text{bwd}}=15$ and of forward solver as $K_{\text{fwd}}=30,50,80$.

Results. We examine results by algorithm and focus on their trend with iteration number (x-axis) and map size (column), not just the absolute numbers. The conclusion already mentioned in the above section: Beyond an intermediate iteration number (around 30-50), IDPs are more favorable because of scalability and computational cost. We present other analyses here.

We start from ConvGPPN and ID-ConvGPPN, which perform the best in ADPs and IDPs class, respectively. They also have the most number of parameters and use greatest time because of the gates in ConvGRU units. As shown in Figure 2, this also caused two issues of ConvGPPN: scalability to larger maps/iterations (out of memory for $27\times 27$ 80 iterations and $49\times 49$ 50 and 80 iterations), and also convergence stability (e.g. $27\times 27$ 50 iterations).

For SymVIN and ID-SymVIN, they replace Conv2D with SteerableConv, with computational cost slightly higher than VIN and much lower than ConvGPPN. Thus, they can successfully run on all tasks and iteration numbers. However, we find that explicit SymVIN may diverge due to bad initialization, and this is more severe if the network is deeper (more iterations), as in Figure 2’s 50 and 80 iterations. Nevertheless, ID-SymVIN alleviates this issue, since implicit differentiable planning decouples forward and backward pass, so the gradient computation is not affected by forward pass.

Furthermore, VIN and ID-VIN are surprisingly less affected by the number of iterations and problem scale. We find their forward passes tends to converge faster to reasonable equilibria, thus further increasing iteration does not help, nor break convergence as long as memory is sufficient.

Forward and backward runtime. We visualize the runtime of IDPs and ADPs in Figure 6. For IDPs, we use the forward-iteration solver for the forward pass and Anderson solver for the backward pass. Note that in the bottom left, we intentionally plot backward runtime vs. forward pass iterations. This emphasizes that IDPs decouple forward and backward passes because the backward runtime does not rely on forward pass iterations. Instead, for ADPs, value iteration is done by network layers, thus the backward pass is coupled: the runtime increases with depth and some runs failed due to limited memory (11GB, see missing dots).

Therefore, this set of figures shows better scalability of IDPs (no missing dots – out of memory – and constant backward time). In terms of absolute time, the forward runtime of IDPs when using the forward solver is comparable with successful ADPs.

Setup. Beyond evaluating generalization to novel maps, we compare their training efficiency with learning curves. Each learning curve is aggregated over 5 seeds, which are from the models in the above section. The learning curves are for all planners on $15\times 15$ maps (Figure 7 left) and $49\times 49$ maps (Figure 7 right, 30 layers for ADPs – due to scalability issue – and 80 iterations for IDPs).

Results. On $15\times 15$ maps, we show $K_{\text{layer}}=80$ layers for ADPs and $K_{\text{fwd}}=80$ iterations for IDPs. ID-ConvGPPN performs the best and is much more stable than its ADP counterpart ConvGPPN. ID-SymVIN learns reliably, while SymVIN fails to converge due to instability from 80 layers. ID-VIN and VIN are comparable throughout training. On $49\times 49$ maps, we visualize $K_{\text{layer}}=30$ layers for ADPs (due to their limited scalability) and $K_{\text{fwd}}=80$ iterations for IDPs. ConvGPPN cannot run at all even for only 30 layers, while ID-ConvGPPN still reaches a near-perfect success rate. ID-SymVIN learns slightly better than SymVIN and reaches higher asymptotic performance. ID-VIN has a similar trend to VIN, but performs worse overall due to the complexity of the task.

Setup. We run all planners on the other three challenging tasks. For visual navigation, we randomly generate 10K/2K/2K maps using the same strategy as 2D navigation and then render four egocentric panoramic views for each location from produced 3D environments with Gym-MiniWorld (Chevalier-Boisvert, 2018). For configuration-space manipulation and workspace manipulation, we randomly generate 10K/2K/2K tasks with 0 to 5 obstacles in workspace. In configuration-space manipulation, we manually convert each task into a $18\times 18$ or $36\times 36$ map ($20^{\circ}$ or $10^{\circ}$ per bin). The workspace task additionally needs a mapper network to convert the $96\times 96$ workspace (image of obstacles) to an $18\times 18$ 2-DOF configuration space (2D occupancy grid). We provide additional details in the Section D.

Results. In Table 1, due to space limitations, we average over $K_{\text{layer}}=30,50,80$ for ADPs and $K_{\text{fwd}}=30,50,80$ for IDPs. For each task, we present the mean and standard deviation over 5 seeds times three hyperparameters and provide the separated results to Section E. We italicize entries for runs with at least one diverged trial (any success rate $<20\%$).

Generally, IDPs perform much more stably. On $18\times 18$ or $36\times 36$ configuration-space manipulation, ID-SymVIN and ID-ConvGPPN reach almost perfect results, while ID-VIN has diverged runs on $36\times 36$ (marked in italic). SymVIN and ConvGPPN are more unstable, while VIN even outperforms them and is also better than ID-VIN. On $18\times 18$ workspace manipulation, because of the difficulty of jointly learning maps and potentially planning on inaccurate maps, most numbers are worse than in configuration-space. ID-ConvGPPN still performs the best, and other methods are comparable. For $15\times 15$ visual navigation, it needs to learn a mapper from panoramic images and is more challenging. ID-ConvGPPN is still the best. ID-SymVIN exhibits some failed runs and gets underperformed by SymVIN in these seeds, and ID-VIN is comparable with VIN.

Across all tasks, the results confirm the superiority of scalability and convergence stability of IDPs and demonstrate the ability of jointly training mappers (with algorithmic differentiation for this layer) even when using implicit differentiation for planners.