Approximate Inference for Stochastic Planning in Factored Spaces

The connection between planning and probabilistic inference is well known and multiple reductions exist showing how inference algorithms can be used to solve stochastic planning problems. Such reductions are equivalent when one can perform exact inference but this is not typically the case for challenging planning problems that have many state variables, a.k.a. factored spaces, where approximate inference schemes are introduced. The planning and reinforcement learning literatures include multiple such efforts where different algorithmic frameworks are combined with different approximation schemes. For example, constructions exist through weighted model counting Domshlak and Hoffmann [2006], several forms of variational inference (e.g., [Toussaint and Storkey, 2006, Levine, 2018]), and several forms belief propagation (e.g., [Liu and Ihler, 2012, Cui et al., 2019]). However, it is not clear how different algorithmic approaches are related to one another and how the choice of approach interacts with the choice of approximation scheme.

The paper makes three contributions. First we provide a unified scheme that connects previous approaches along two dimensions, using either forward or backward reasoning, and choosing what approximation to use, where we address Belief Propagation (BP) and mean field variational inference (MFVI). This allows us to put prior work in a unified framework that explains choices made by corresponding algorithms. In particular, our analysis shows that Forward MFVI which is used in some papers can be understood to run multiple iterations of Backward MFVI and thus provides tighter approximations. Second, through extensive experiments over large planning problems, we show that forward reasoning with Belief Propagation provides the best performance among these algorithms, that MFVI provides poor performance in some domains, and that modifying MFVI using exponentiated rewards helps in some cases but not sufficiently. We also analyze the failures of MFVI experimentally pointing to sensitivity in updates. Third, based on the analysis, we propose a novel algorithm, Collapsed State Variational Inference (CSVI), that uses mean field with collapsed variational inference where state variables are integrated out. CSVI is motivated theoretically due to its tighter variational approximation and we show empirically that it matches the performance of Forward Belief Propagation. This shows that while naive application of mean field for planning fails, other variational approximations like CSVI can yield strong planning performance. Due to space constraints some technical details and experiment results are omitted from the main paper and are provided in [Wu and Khardon, 2022], which we refer to below as the the appendix.

We consider finite horizon MDPs as specifications of planning problems. Such specifications are often compiled from a high level description language but this is orthogonal to the discussion in the paper. Specifically, consider Markov Decision Processes $\langle\mathcal{S},p(s_{0}),\mathcal{A},\mathcal{P},\mathcal{R},T,\gamma\rangle$, where $\mathcal{S}$ denotes the state space, $p(s_{0})$ is a distribution over start states, $\mathcal{A}$ denotes the action space, $\mathcal{P}$ denotes the transition probability $p(s_{t+1}|s_{t},a_{t})$, $\mathcal{R}$ denotes the reward function $R(s_{t},a_{t})$, $T$ denotes the horizon and $\gamma$ is the discount factor. In this paper we set $\gamma=1$, but this does not significantly affect any of the formulations. A solution is given by a policy, $\pi$, that specifies $p_{\theta_{t}}(a_{t}|s_{t})$ with policy parameters $\theta=\{\theta_{t}\}$ allowing for non-stationary policies. The task in planning is to find a policy that maximizes the expected cumulative reward $\mathbb{E}[\sum_{t=0}^{T-1}R(s_{t},a_{t})]$ where the expectation is taken w.r.t. trajectories generated from the MDP with policy $\pi$, that is, $s_{0}\sim p(s_{0})$, $a_{t}\sim p_{\theta_{t}}(a_{t}|s_{t})$, and $s_{t}\sim p(s_{t}|s_{t-1},a_{t-1})$.

In this paper we follow recent practice in stochastic planning and use the online planning framework, where in a state $s$, the algorithm computes for a limited time to pick an action $a$, uses $a$ to control the MDP to get to the next state, and repeats this process. Online planning is often used with receding horizon control, where the planner uses a $T$ step lookahead in its search and then extracts the first action $a$ to be applied in $s$.

Solving a finite horizon MDP is equivalent to solving an inference problem in the corresponding Dynamic Bayesian Network (DBN), or more precisely in the dynamic decision network. We assume a factored form of states consisting of binary state variables, i.e. $s_{t}=(s_{t}^{1},\cdots,s_{t}^{M})$. We also assume a factored action representation $a_{t}=(a_{t}^{1},\cdots,a_{t}^{N})$.

The MDP formulation above requires real-value reward nodes in the DBN. To facilitate inference one can replace these nodes with constructions that use only binary variables, and various such constructions appear in the literature. In the following we develop one such construction and use that in our experiments. We introduce binary reward random variables $r_{t}$ to capture the reward after taking action in the previous time step $t-1$, the distribution of which is defined as

We can then define $\tilde{R}$ where $p(\tilde{R}=1)=\frac{\sum_{1}^{T}r_{t}}{T}$ to capture the cumulative reward. We call the resulting DBN the intermediate representation. However, $\tilde{R}$ has $T$ parents which hinders efficient inference. To avoid the use of $\tilde{R}$, we introduce cumulative reward binary random variables $c_{t}$. To keep the consistency of the graphical structure, we create an auxiliary node $c_{0}\equiv 1$, and for $t>0$ the distribution of $c_{t}$ is defined recursively depending on the previous cumulative reward $c_{t-1}$ and current reward $r_{t}$:

In many planning problems the reward is given as an additive function over a set of small factors. For such problems we introduce another chain of binary reward variables within a time step using a similar construction. This yields a DBN that only includes binary variables with a small number of parents. As the following proposition shows the three constructions, using cumulative reward, using $\tilde{R}$ and using $c_{T}$ are equivalent. Further details and proofs are given in the the appendix.

In the following we restrict our discussion to open loop policies, that is, $p_{\theta_{t}}(a_{t}|s_{t})=p_{\theta_{t}}(a_{t})$ where the policy is time dependent but does not depend on the state (other than $s_{0}$ if it is fixed). Thus for an action sequence $A=\{a_{0},\ldots,a_{T-1}\}$, we have $p_{\theta}(A)=\prod p_{\theta_{t}}(a_{t})$. This covers most of previous work on planning as inference in the literature. The extension to standard policies is straightforward but requires more complex algorithms for optimization.

The idea of using inference for stochastic planning has a long history and has attracted many different approaches. For example, Cooper [1988] showed how inference can be used for decision making in influence diagrams, Domshlak and Hoffmann [2006] use an approach based on weighted model counting, Nitti et al. [2015] use a probabilistic programming formulation, and Lee et al. [2021] use anytime marginal MAP solvers for planning problems.

Several groups have developed approaches that follow the forward variational framework, going back to Dayan and Hinton [1997]. This idea is often developed by defining a reward weighted path distribution which is similar to conditioning on $c_{T}=1$ in our framework, and developing algorithms from this formulation [Furmston and Barber, 2010, 2011, Toussaint and Storkey, 2006, Kumar et al., 2015]. We note, however, that these works did not explicitly address factoring over state and action variables.

On the other hand, some papers in robotics and reinforcement learning (RL) [Toussaint, 2009, Kappen et al., 2012, Levine, 2018] follow the backward variational framework. In contrast with the discussion above they use a formulation where the reward over trajectories is exponentiated. As shown by Levine [2018] this modifies the original optimization objective by adding a term with the expected entropy of the policy, and hence solves a slightly different problem, but the entropy term may be beneficial for exploration in RL. In addition, the work of Neumann [2011] uses the forward variational algorithm, but with an exponentiated reward, and additional sampling-based approximations. We can see that the forward and backward variational approaches have been widely used but have not been differentiated before. Our analysis above clarifies the relationship between these approaches.

For the case of BP approximation, Murphy and Weiss [2001] proposed the Factored Frontier Algorithm which is a forward BP method for marginal inference, and Boyen and Koller [1998] developed approximation bounds for forward inference. The work of Liu and Ihler [2012], Kiselev and Poupart [2014] follows the forward BP framework, but develops a generalized belief propagation algorithm that solves both optimization and expectation steps using message passing. The work of Cui et al. [2019] also follows the forward BP framework but decouples the expectation which is done through BP from the optimization that uses an approximation based on gradient search.

Several works have made additional assumptions on the structure of the DBN in their discussion of graph-based MDPs. Cheng et al. [2013] extend the algorithm of Liu and Ihler [2012] to this case. Peyrard and Sabbadin [2006] and Sabbadin et al. [2012] use the Mean Field approximation method but only use it to approximate the distribution over state variables. They then use the approximate distribution to approximate steps of the Policy Iteration algorithm. Hence their algorithm is different from MFVI in that reward variables are not included in the variational approximation. Finally, our work can be seen to extend the comparison of Mean Field and Loopy BP for general inference tasks [Weiss, 2001]. As in this early work, our experiments show that optimization of variational objectives can lead to local optima and that BP can provide some advantage.

This section presents an experimental evaluation of the algorithms. The code for regenerating all the results is available on Github¹¹1https://github.com/Zhennan-Wu/AISPFS. Our goal in this paper is to understand the quality of decisions provided by different approximate inference schemes, ignoring implementation details. Therefore, during the experiments we do not limit run time but instead allow the algorithms to converge, within bounds given below, before proposing a decision. We chose 6 problem domains from the ICAPS 2011 International Probabilistic Planning Competition to conduct our experiments. Each domain has 10 instances with factorized structure, horizon of 40 and discount factor of 1, and instances differ by the number of state and action variables. For our experiments we use the SPUDD [Hoey et al., 1999] translation of the original RDDL [Sanner et al., 2010] specification, which compiles away action factoring. This simplifies the implementation because it removes the need to reconcile action constraints with factoring. To control our overall experimental time we use online planning with receding horizon control, where we set the search horizon to be the minimum value between 9 and the remaining time steps.

Algorithmic parameters for MFVI: we perform at most 100 Variational updates and stop early if the infinity norm of the difference between consecutive approximation distributions is less than 0.1. We perform 3 outer iterations, i.e., policy updates for the forward version.

Algorithmic parameters for BP variants: We use SOGBOFA [Cui et al., 2019]²²2https://github.com/hcui01/SOGBOFA as forward Loopy BP, fixing search depth to 9, and limiting the number of gradient updates to 500. We note that SOGBOFA has outperformed other planners, including search based planners, in IPPC 2018 problems and is a state of the art baseline for the evaluation. For the backward algorithm, our implementation is based on Zhou et al. [2022] with parallel message update and a bound of 100 iterations with no damping ($\beta=0$).

Normalized mean $\pm$ one standard deviation of the cumulative reward over 12 simulations are shown in all the plots. Denote the mean value and standard deviation of the cumulative reward of algorithm $a$ on instance $i$ to be $\bar{r}_{i}(a)$, $\sigma_{i}(a)$, respectively. To facilitate comparisons across domains we report scores normalized relative to the random policy. Specifically, for algorithm $a$ on instance $i$, $\mbox{score-mean}_{i}(a)=\frac{|\bar{r}_{i}(a)-\bar{r}_{i}(RANDOM)|}{|\bar{r}_{i}(RANDOM)|}$ and $\mbox{score-std}_{i}(a)=\frac{\sigma_{i}(a)}{|\bar{r}_{i}(RANDOM)|}$ where the random algorithm has score 0 and higher scores indicates better performance. For reference, the raw results are given in the the appendix.

Motivated by the analysis above, we propose a new algorithm for variational inference in planning. Instead of treating all the latent nodes in the DBN in the same manner and computing approximate distributions over all these variables, the algorithm focuses on the action variables and effectively marginalize out other terms to achieve a tighter ELBO. This type of approach is known as collapsed variational inference, which has been shown to be effective in models where the marginalization can be done analytically (e.g., Teh et al. [2006]) but for planning one has to resolve additional computational challenges as we show below. Specifically we propose to use the following provably tighter ELBO

Here we have the same factorized transitions and policy distribution. However, we do not compute approximation distributions over state, reward, and cumulative reward variables. With mean field, the standard solution [Bishop, 2006] yields the update equation

The tighter approximation appears to yield an infeasible update, because $A$ is entangled in $g()$ and we must perform an explicit marginalization in $g()$ for each update.

We next show how the update equation can be approximated via sampling. The key is to first extract $p_{\theta}(A)$ from the expectation. We therefore have:

Recall that $p_{\theta}(A)$ is a product of independent terms. This implies that the first part can be substituted with $\log p_{\theta}(a^{l}_{h})$ since all other terms are constants w.r.t the variable of interest in (7) and they will vanish in the normalized update of $q_{\phi}(a_{t}^{l})$. The second part is conditioned on $A$ and does not include $p(A)$ terms. Its expectation can be estimated through sampling. In particular, sampling can be intuitively done as follows: keeping $a^{l}_{t}$ fixed, sample the action sequence from approximate distribution $q_{\phi\backslash a^{l}_{t}}(A)$. Then complement this by sampling values for $s_{t}$, $r_{t}$, $c_{t}$ nodes, including $c_{T}$. The resulting values for $c_{T}$ are generated from the correct distribution and the average over $c_{T}$ gives an estimate of the expectation. Since we are using sampling and averaging inside the logarithm this yields biased estimates for updates, but this type of biased estimates has been shown to work in other cases in machine learning (e.g., [Wei et al., 2021]) and it can be mitigated by taking sufficient samples. It is interesting to note from the above update that the policy distribution serves as a weight bias in the action update procedure. Algorithm 3 summarizes the update procedure.

In this paper we provide a unified scheme that categorizes many previous approaches along two dimensions, using either forward or backward reasoning and choosing an approximation scheme. Specifically, we focus on belief propagation and mean field variational inference as the approximation choices. In this context, we illustrate the advantage of Forward Loopy BP as providing the best performance. Algorithms based on MFVI perform poorly in some domains. They are improved by exponential reward weighting but not sufficiently so. An experimental analysis points to sensitivity of the optimization as a source for this failure. Motivated by this analysis we propose a novel algorithm, Collapsed State Variational Inference, which provides a tighter variational approximation, and while being computationally demanding it performs competitively with Forward Loopy BP. The results highlight that while BP has been less in focus in recent years, it provides a strong baseline for stochastic planning. It also shows the importance of focusing variational approximations on variables of interest as done in CSVI and the potential for developing strong variational algorithms for planning. These observations suggest interesting directions for future work including developing efficient variants of CSVI, using amortized variational inference in planning to improve CSVI, alternative schemes to capture the posterior distributions in VI, and developing tighter approximations and optimization algorithms through BP methods.

This work was supported by NSF under grant IIS-1906694 and grant IIS-2002393. Some of the experiments in this paper were run on the Big Red 3 computing system at Indiana University, supported in part by Lilly Endowment,Inc., through its support for the Indiana University Pervasive Technology Institute.