DiSProD: Differentiable Symbolic Propagation of Distributions for Planning

1 Introduction

Planning is one of the key problems in artificial intelligence (AI), as it enables intelligent agents to make informed decisions and achieve their objectives in complex and dynamic environments. This is especially important when the environment is inherently stochastic or when the dynamics model used by the planning algorithm is imperfect. As a result, research on planning with stochastic transitions has been multifaceted, encompassing symbolic task-level planning in discrete spaces [Kolobov et al., 2012; Keller and Helmert, 2013; Cui et al., 2019], robotic motion planning in continuous spaces [Kurniawati et al., 2011; Van Den Berg et al., 2012; Agha-Mohammadi et al., 2014], integrated task and motion planning [Kaelbling and Lozano-P´erez, 2013; Vega-Brown and Roy, 2018; Garrett et al., 2021] and model-based reinforcement learning [Chua et al., 2018; Hafner et al., 2020; Curi et al., 2020].

Markov Decision Processes (MDPs) provide the theoretical foundation for planning under uncertainty, but scalability remains a challenge, and approximation is often necessary.

Many approaches have been proposed in the literature, including searching in action space, searching in state space, approaches using Monte-Carlo simulation, and approaches using differentiation. Several gradient-based planners have been proposed, but they are not easily usable as domain independent planners due to scalability [Deisenroth and Rasmussen, 2011], or restriction to deterministic environments [Wu et al., 2020a], or in the Reinforcement Learning (RL) context, where, in most works the success of the planner depends on the quality of the learned domain-specific value function [Tamar et al., 2016; Hafner et al., 2020].

The paper fills this gap by introducing a novel domainindependent online planner using differentiable probabilistic transition models. Our work is related to prior algorithms for trajectory optimization, distribution propagation, and differential dynamic programming [Chua et al., 2018; Williams et al., 2017; Deisenroth et al., 2013; Tassa et al., 2012; Lenz et al., 2015]. However, we introduce a novel symbolic approximation and propagation scheme generalizing work on AI planning that facilitate robustness, better approximation, and optimization [Cui et al., 2019].

The main contribution of this work is in designing DiSProD (Differentiable Symbolic Propagation of Distributions), an online planner for environments with probabilistic transitions, in continuous or hybrid spaces. The core idea is to create a symbolic graph that encodes the distribution of future trajectories conditioned on a given policy. The resulting symbolic graph provides an analytically differentiable representation of the policy's value, allowing efficient gradient-based optimization of the action variables for long-horizon search. While distributions over trajectories are too complex to be captured exactly, DiSProD uses Taylor's approximation and an independence assumption to facilitate a symbolic propagation of product distributions over future state and reward variables. The approximate distribution propagation can be viewed as an efficient symbolic aggregation of many trajectories, which differs from sampling algorithms that aggregate the results of many individual trajectories. This approach reduces variance in estimates in stochastic environments and facilitates planning with sparse rewards.

Extensive quantitative experiments are conducted to compare DiSProD with state-of-the-art planners in discrete-time planning in OpenAI Gym environments [Brockman et al., 2016] and continuous-time control of simulated robotic systems. The results show that DiSProD outperforms existing planners in dealing with stochastic environments, sensitivity to search depth, sparsity of rewards, and large action spaces.

Furthermore, we use DiSProD with an approximate transition model to control two real-world robotic systems demonstrating that it can successfully control ground vehicles and surface vessels to navigate around obstacles.

Due to space constraints, some details are omitted from the paper. The full paper as well as code to reproduce the experiments and videos from physical experiments are available at https://pecey.github.io/DiSProD.

2 RelatedWork

Planning in continuous spaces has been studied in several sub-fields. A key distinction in planning methods is between offline planners, which compute a complete solution and then apply it, and online planners or Model Predictive Control (MPC) [Borrelli et al., 2017], in which at every timestep, optimization is carried out over a finite horizon and only the first action from the solution is executed. The online nature of MPC provides some robustness to unexpected outcomes at the cost of increased computation time during action selection. Another important distinction is planning in state space (or configuration space) [LaValle, 2006] versus planning in action space. The former seeks an optimal sequence of states, leaving action execution to a low-level controller, while the latter produces executable actions directly. DiSProD is an online action planner but it can also produce a sequence of states as a byproduct, as we discuss in Section 4.6.

Within online action planners, DiSProD is related to two lines of work using differentiable transition models. First, our approach builds on the SOGBOFA algorithm [Cui et al., 2018; Cui et al., 2019], which was developed for discrete task-level AI planning problems. However, that work is restricted to binary state and action variables. In addition, DiSProD introduces a new distribution propagation method based on symbolic Taylor expansions. The second group includes planners using differentiable transition models in RL and control. However, deep RL work (e.g., [Heess et al., 2015; Depeweg et al., 2017]) makes use of learned value functions to aid planning performance and cannot plan in a new model without training first, and many approaches (e.g., [Wu et al., 2020a]) use deterministic transition models. In addition, most approaches [Hafner et al., 2020; Bueno et al., 2019] optimize over individual trajectories and do not propagate distributions over trajectories as in DiSProD. In this realm, iLQG [Tassa et al., 2012] and the PILCO family [Deisenroth et al., 2013; Parmas et al., 2018; Kamthe and Deisenroth, 2018] are most related to our method. iLQG linearizes the dynamics, assumes linear Gaussian transitions, as in the Extended Kalman Filter (EKF), and optimizes over individual trajectories. PILCO does propagate distributions analytically, albeit with the restricted Gaussian process (GP) dynamics and the Gaussian kernel. Gal et al. [2016] replace the GP in PILCO with Bayesian neural networks, which cannot propagate distributions analytically and hence requires particle-based planning. From this perspective, DiSProD can be seen as a generalization of PILCO and iLQG, which uses differentiation over approximate symbolic propagation of distribution.

Our work is also related to sampling-based planners Model Predictive Path Integral (MPPI) and Cross-Entropy Method (CEM) [Kobilarov, 2012; Williams et al., 2017; Chua et al., 2018; Wagener et al., 2019; Mohamed et al., 2022]. These algorithms sample a set of trajectories around a nominal trajectory using a fixed sampling distribution, and update the nominal trajectory based on the "goodness" of the trajectories to bias sampling in subsequent iterations. Similarly, Mania et al. [2018] use sampling to estimate numerical gradients which provide related policy updates. Our framework uses distribution propagation instead of sampling and it optimizes stochastic policies instead of having a fixed sampling distribution.

Finally, we note that our approach can flexibly handle both discrete and continuous variables. In contrast, other methods for planning in hybrid spaces typically restrict the transitions, for example, to piecewise linear models [Li and Littman, 2005; Zamani et al., 2012; Raghavan et al., 2017].

3 Algorithms andMethodology

A Markov Decision Process (MDP) is specified by {S,A, T,R, γ}, where S is the state space, A is the action space, T is the transition function, R is the one-step reward function, and γ is the discount factor. A policy π is a mapping from states to actions. Given a policy π, the action-value function Qπ (st, at) represents the expected discounted total reward that can be obtained from state st if action at is chosen, and the policy π is followed thereafter, where st and at are vectors of state and action variables at time t. We focus on open loop probabilistic policies parameterized by θ = {θt} where each θt picks the action at time step t and is further factorized over individual action variables. As in prior work, we approximate Qθ (st, θt) = E[∑D-1 i=0 γ iR(st+i, at+i)] where at ∼ p(at|θt) and st+1 ∼ p(st+1|st, at), optimize θ and pick the action using θt. In principle the Q-value can be calculated from the distributions over {(st+i, at+i)} but these distributions are complex. We approximate these distributions as products of independent distributions over individual state variables.

Our algorithm can be conceptually divided into two parts. The first calculates an approximate distribution over future states, rewards and Qπ (st, at), conditioned on a given policy. A schematic overview of this process is shown in Figure 1. The second uses this representation to optimize the policy. These are described in the next two subsections.

3.1 Analytic Computation Graph

Transition Model. Typically, given st and at, a simulator Tsim, samples from a distribution over st+1. Formally, st+1 ∼ Tsim(st, at). For DiSProD, we need to encapsulate any internal sampling in Tsim as an input, to separate sampling from computation, similar to the reparameterization trick.

Tsim and its encapsulation are shown in Figure 1a and 1b.

Representation. Let st,k denote the value of a randomvariable sk at time t. If st,k is binary, then its marginal distribution can be captured using just its mean ˆµs,t,k as the variance is implicitly given by ˆ vs,t,k = ˆµs,t,k(1 - ˆµs,t,k). If st,k is continuous, its distribution can be summarized using its mean and variance (ˆµs,t,k, ˆ vs,t,k). Similarly, the distribution over the noise variable ϵi and action variable aℓ are represented using (ˆµϵ,t,i, ˆ vϵ,t,i) and (ˆµa,t,ℓ, ˆ va,t,ℓ) respectively. While the use of mean and variance suggests a normal distribution, the proposed construction does not assume anything about the form of the distribution.

Independence Assumption. To simplify the computations of propagated distributions, our approximation assumes that for all t, the distribution over the trajectories given by p(st, at, ϵt) is a product over independent factors: p(st, at, ϵt) = ∏k p(st,k) ∏ℓ p(at,ℓ) ∏i p(ϵt,i).

Approximation. Our main observation is that the above assumption suffices to propagate approximate distributions and approximate the Q-function. Specifically, we approximate the transition function using a second order Taylor expansion whose terms are given analytically in a symbolic form. Let zt = (st, at, ϵt), ˆ zt = (ˆµst , ˆµat , ˆµϵt ) and st+1 = T(zt). To simplify the notation, consider the j-th state variable st+1,j = Tj(st, at, ϵt) = Tj(zt) and let ∇Tj = ∂Tj/∂zt and Hj = ∂ 2 Tj/∂zt∂z⊤ t . We use Taylor's expansion to approximate the encapsulated model in Equation (1) by

This is illustrated in Figure 1c. We use this approximation to calculate the mean and variance of st+1 via methods of propagating distributions. Given our independence assumption, the off diagonal covariance terms multiplying Hj are zero, and the expected value of the st+1,j becomes

A similar approximation for the variance yields 4th order moments. To reduce complexity, we use a first order Taylor approximation for the variance resulting in

To write these concisely we collect the first order partials and diagonals of the Hessians into matrices as follows:

Similarly, we define J T at , ˜HT at , J T ϵt , and ˜HT ϵt for action and noise variables, respectively. We also define ˜HR st , ˜HR at , and

˜HR ϵt to be the second order partials of the reward function R with respect to the state, action and noise variables. We can now write the vector form of the equations as follows:

In this paper, we do not model the variance of the reward function. However, this can be easily done by analogy with Equation (8) which will facilitate risk sensitive optimization, for example, conditional value at risk [Chow et al., 2017]. The computations of Equations (7) to (9) are illustrated in Figure 1d. Stacking these computations over multiple time steps gives us an approximation of the distribution over future states, captured analytically as a computation graph, as shown in Figure 1e. The Q-function is then approximated as

where we use γ = 1 in our experiments. Note that the computation graph propagates distributions and does not sample trajectories. The computation only requires the mean (ˆµϵ = 0) and variance (ˆ vϵ = 1) of the noise which are known in advance and are absorbed as constants in the graph.

Algorithm 1 One Step of DiSProD.

Input: state

• 1: initialize actions (for all restarts)
• 2: build computation graph till depth D
• 3: while actions have not converged do
• 4: loss = - ∑restart k ˆQ(st, ˆµka, ˆ vka)
• 5: loss.backward()
• 6: actions ← safe-projected-gradient-update(actions)
• 7: save action-means ˆµk∗ at+1:t+D from the best restart k∗
• 8: return action ∼N(ˆµk∗ at , ˆ vk∗ at )

Our construction is very general - we only require access to T(st, at, ϵt), R and analytic computation of their partial derivatives. In practice, T and R can have non-differentiable components which can be mitigated by approximating the non-smooth functions with their smoother alternatives.

3.2 Optimization Algorithm

Thanks to the symbolic representation, once the computation graph is built, it can be reused for all steps of problem solving as long as the reward function does not change. The optimization is conceptually simple, initializing all action variables for all time steps and performing gradient based search. However, some important details are discussed next.

Multiple Restarts and Loss function. Following Cui et al. [2019], we perform gradient search with multiple restarts. While they performed restarts sequentially, we take advantage of the structure of the computation to perform the restarts in a vectorized form. Recall that each restart is represented by (ˆµa, ˆ va) as in Equation (10). Using a superscript k to represent independent restarts we can write the loss function (negating Q) as loss = - ∑restart k ˆQ(st, ˆµka, ˆ vka). Now to evaluate the loss we generate a matrix where every row represents a possible policy (ˆµka, ˆ vka) and evaluate ˆQ on this matrix. Therefore, we benefit from evaluating all restarts as a "batch of examples" relative to the computation graph. Since we are optimizing with respect to the input matrix (and not computation graph weights), the loss decomposes into a sum of independent functions and gradient search effectively optimizes all restarts simultaneously.

Initialization of Actions. For discrete action variables, following Cui et al. [2019], we initialize actions parameters for time step t = 0 to binary values and for steps t > 0, we initialize the marginals to random values in [0, 1]. For continuous action variables, we must constrain actions to be in a valid interval. For each action variable aℓ, the policy is given by a uniform distribution, initializing ˆµa,ℓ ∼ U(aℓ,min, aℓ,max) and the variance of a uniform distribution centered around ˆµa,ℓ, i.e., ˆ va,ℓ = min(aℓ,max-ˆµa,ℓ,ˆµa,ℓ-aℓ,min)2/12.

Policy Updates. We use Adam [Kingma and Ba, 2015] to optimize the action variables over all restarts simultaneously. We make at most K = 10 updates and stop the search early when ||ˆµnew a - ˆµold a ||∞ ≤ 0.1 and ||ˆ vnew a - ˆ vold a ||∞ ≤ 0.01 for normalized action ranges. Gradients are used to update both ˆµa and ˆ va, which implies that we search over a stochastic policy. This can be important for success in sparse-reward scenarios where a stochastic policy effectively broadens the search when needed, while also allowing a nearly deterministic choice of actions at convergence.

Safe Projected-Gradient Update. The gradient-based updates have to satisfy two constraints. The first is the need to constrain variables into valid intervals. Similar to [Tassa et al., 2014; Cui et al., 2019], we use the standard projected gradient descent to restrict ˆµa between amin and amax, while ˆ va is constrained to min(1/12, min(ˆµa-amin,amax-ˆµa)2/12) which is the variance of the largest legal uniform distribution centered around ˆµa. If the gradient step pushes ˆµa or ˆ va outside the valid region, it is clipped to the boundary. Finally, we only take a gradient step on the restart if it improves the Q-value and otherwise we maintain the previous value. This safe update requires an additional evaluation of the computation graph to check for the Q-value improvement and it increases runtime, but it ensures that ˆQ is monotonically increasing.

Saving Actions. Since the gradients are back-propagated through the entire computation graph, action variables at all depths are adjusted according to the corresponding gradients. However, only the action at depth d = 0 is executed. Formally, at time t, we sample an action using ˆµat and ˆ vat while ˆµat+1:t+D and ˆ vat+1:t+D are not used. These updated actions can be used to initialize the action variables when planning at state st+1. The same idea has been used before in MPPI [Williams et al., 2017; Wagener et al., 2019]. Note that MPPI also uses multiple samples, but these samples all contribute to the update of a single action sequence, which can potentially harm the search [Lambert et al., 2021; Barcelos et al., 2021]. This is different in DiSProD where each restart is an independent search. To allow reuse of old search but add diversity, we initialize one restart using saved action mean, and initialize all other restarts randomly.

Overall OptimizationAlgorithm. These steps are summarized in Algorithm 1, where search depth (D), the initial stepsize for ˆµa and ˆ va (lrµ and lrv), and number of restarts serve as hyper-parameters. When all restarts have converged or K gradient steps have been performed, we choose the best restart with the maximum Q-value, breaking ties randomly. Finally, we sample an action using the mean and variance of the first action distribution of the best restart.

3.3 Discussion

DiSProD propagates approximate distributions over trajectories using moment matching on individual variables, assuming independence between variables. As shown by Cui et al. [2018], in the binary case, this is equivalent to belief propagation, which works well empirically despite lack of formal guarantees. DiSProD offers a different trade-off from sampling-based methods which are exact in the limit of infinite samples but are sensitive to variance in estimates with limited samples. In this sense, our computation graph provides a stable if biased estimate of the Q-function. With deterministic transitions and policy, our computation is exact and gradient descent can potentially find the optimal policy. For stochastic transitions, the computation is approximate and we use a loose stopping criterion to reduce run time.

4 Experiments

We experiment with DiSProD on a variety of deterministic and stochastic environments and evaluate its robustness to long horizon planning, its ability to work with sparse rewards, and its performance in high-dimensional action spaces. For this purpose, we conduct extensive experiments on simulated and real-world environments with increasing complexity.

We use three OpenAI Gym environments (Cart Pole, Mountain Car, and Pendulum) to evaluate the robustness of the compared planners in terms of stochasticity, planning horizon, and reward sparsity. The original OpenAI Gym environments are deterministic. We therefore enhance these environments by explicitly adding noise into the model and using it as a part of the dynamics T(s, a, ϵ). Further details are in the full paper. For the discussion, it suffices to note that we parameterize the amount of noise using a scalar α and can therefore evaluate planners' performance as a function of stochasticity in the dynamics.

In addition, we developed a new Gym environment that models a simplified vehicle dynamics, elaborating over the Dubins' car model [Dubins, 1957]. In particular, the agent can control the change in linear (∆vt) and angular velocity (∆ωt) instead of controlling the velocities directly, and the maximum change (∆v,∆ω) is limited to a small value. This model is still inaccurate because it ignores friction, inertia, actuation noise and localization noise, but it provides an acceptable approximation. We then use this model to plan in a physics simulator with asynchronous execution. Specifically, the planner is used to control a TurtleBot in Gazebo simulation via a Robot Operating System (ROS) interface.

Finally, to demonstrate robustness and applicability, DiSProD is used to control two physical robot systems: an Unmanned Ground Vehicle (Jackal) and an Unmanned Surface Vessel (Heron) using the aforementioned model.

In the experiments we use 200 restarts, and for DiSProD lrv = lrµ/10. The values of D and lrµ for Gym environments are provided in Figure 2. For TurtleBot we use D = 100, lrµ = 10. For Jackal and Heron, D is modified to 30, 70 respectively, and we use 400 restarts.

4.1 Baselines

We compare DiSProD to CEM and MPPI. Both are shootingbased planners that maintain a sequence of action distributions to sample actions from. From a given state, they use the sampled actions to generate multiple trajectories and compute the rewards accumulated at the end of each trajectory. The two algorithms differ in how they recompute the action distribution. While CEM uses the actions from the the top n trajectories to form the new action distribution, MPPI weights the actions in a trajectory by the cumulative reward for that trajectory. Hyperparameters for the planners were frozen after tuning them on deterministic versions of the environments.

CEM and MPPI can potentially benefit from the use of saving actions. We have found that this is helpful in the basic environments but harms their performance in the car model. To ensure a fair comparison, we use the best setting for the baselines whereas DiSProD always saves actions.

4.2 Evaluation in Basic GymEnvironments

Increasing Stochasticity. First, we explore the performance of the planners with the original deterministic environments and with added stochasticity. To separate randomness in the environment from randomness in the experiment, we perform 8 repetitions over 6 runs in each environment, calculating means in each repetition and standard deviations of the means across repetitions. Averages and standard deviations are shown in Figures 2a to 2c. We observe that the planners perform similarly in deterministic environments (α = 0) but with increasing amounts of stochasticity, DiSProD degrades more gracefully and performs better than CEM or MPPI.

Increasing Planning Horizon. Intuitively the deeper the search, the more informative it is. But for sampling-based planners, deeper search can also increase variance when estimating action quality and hence harm performance. Results for experiments testing this aspect in noisy variants of basic environments are shown in Figures 2d to 2f and deterministic environments are included in the full paper. We observe that while in the deterministic setting all the planners have similar performance, in some stochastic environments, the performance of CEM and MPPI can degrade with increasing search depth while DiSProD gives better results. In these environments, beyond a required minimum depth, increasing the depth further does not help performance. This changes, however, when we change reward sparsity.

Increasing Reward Sparsity. Robust planners should be able to work with both sparse and dense rewards. Intuitively, if the reward is dense then even a short search horizon can yield good performance, but a deeper search is required for sparse rewards. To test this, we evaluate the performance of the planners by varying the sparsity of the rewards in the Mountain Car environment. With the standard reward function, the agent gets a large positive reward when its position and velocity are larger than certain thresholds. We modify the reward function to use a smooth version of greaterthan-equal-to function given by σ(10β(x - target)) where β is a sparsity multiplier (β = 1 in earlier experiments) and σ(a) = 1/(1 + e-a ). The larger the value of β, the harder it is to get a reward from bad trajectories. Results are shown in Figures 2g and 2h. Figure 2g shows that with the standard search depth of 100, all planners fail to reach the goal once the reward becomes very sparse. Figure 2h shows that DiSProD can recover and perform well by increasing search depth to 200 but CEM and MPPI fail to do so.

4.3 Evaluation in High Dimensional Action Space

To explore high dimensional action spaces, we modify the Mountain Car environment by adding redundant action variables. The dynamics still depend on a single relevant action variable but the reward function includes additional penalties for the redundant actions. To obtain a high score, the agent must use a similar policy as before for the relevant action and keep the values of the redundant action variables as close to 0 as possible. Details of the model are in the full paper. We compare the performance of the planners against number of redundant actions, without changing any other hyperparameters. Results are shown in Figure 2i, where the MPPI results are omitted since the scores are too low and they distort the plot. We observe that while CEM/MPPI perform poorly as the action space increases, the performance of DiSProD remains stable. The performance of CEM improves if we increase the population size (number of samples) by a factor of 100 (from 200 to 20,000), but it still lags behind DiSProD . Additional results analyzing this scenario are in the full paper. We note that despite the inferior reward, CEM is able to reach the goal location. However, as shown in Figure 2j, it requires a lot more steps to reach the goal. This experiment illustrates the potential advantage of planners that use the analytic model to identify what causes good outcomes as compared to estimating this effect through sampling.

4.4 Evaluation with Hybrid State Space

While we focus on experiments in continuous spaces, our planner is compatible with hybrid environments. To illustrate this we modify the Cart Pole environment to include a binary variable which is set to 1 if the cart is to the right of a certain x-coordinate. The agent receives a reward of 3 when this binary variable is set to 1, otherwise it gets a reward of 1. Figures 2k and 2l show the performance of DiSProD and CEM against planning horizon and noise level. We observe that a deeper search is required for this problem and that DiSProD can successfully obtain high reward.

4.5 Evaluation with a Physics Simulator

We control a TurtleBot using DiSProD on 16 maps with varying degrees of difficulty. In these experiments, the analytic transition model does not take obstacles into account. Incorporating obstacles in the transition yields similar or better results, but slows down the planner due to the increased size of the analytical computation graph. Instead, the reward function penalizes the agent on collision with obstacles. The obstacle patterns and detailed evaluation results are in the full paper. Following Wu et al. [2020b], we use success rate (SR) and success weighted by optimal path length (SL) for evaluation. SR measures the success percentage across different maps, while SL is the ratio of actual path length to the euclidean distance from starting position to goal, averaged over cases where navigation was successful. Intuitively, the lower the SL value, the faster the agent reaches the goal.

Results are shown in Table 1, averaged over 5 runs for each map and averaged over all instances. We observe that all planners are able to control the TurtleBot but DiSProD performs better in both metrics. We also evaluated the planners in a Gym environment where the planners' model matches the environment dynamics exactly. In this case, all planners perform similarly. Hence performance differences in TurtleBot are mainly due to better handling of the inaccurate model.

4.6 Experiments with Robotic Systems

We use DiSProD with the same analytic model to control Jackal and Heron. Real-time control requires high frequency control commands to avoid significant drift. Therefore, one has to find a good balance between planning horizon (D) and the maximum change of linear velocity ∆v. Experiments with low ∆v require a large D and are slow, while experiments with high ∆v suffer from drift, as the hardware cannot stop or accelerate instantaneously. We specified these parameters through initial exploration with the systems. Other parameters are the same as in the TurtleBot simulation.

For experiments with Jackal, DiSProD is used exactly as above, i.e., sending actions to the robot. For the surface vessel, however, we use our planner in another way because our Heron has a motor issue in one of its thrusters. Specifically, DiSProD optimizes the action sequence exactly as before. Instead of sending the actions, it computes intermediate states expected to be reached with the policy (which are available in the computation graph), and sends these as "waypoints" to a PID controller. We found that, since our planner works at a fine time granularity, we can send every 5th state to a PID controller and achieve smooth control. Figure 3 visualizes the trajectories generated by DiSProD when controlling Jackal and Heron. Some videos from these experiments can be seen at https://pecey.github.io/DiSProD.

4.7 Ablation Study and Runtime Comparison

We first evaluate the contribution of the different variance terms (ˆ vst , ˆ vat and ˆ vϵt ) in Equations (7) and (8) to the performance of the algorithm. We use the term complete mode when the planner uses the variance terms (DiSProD) and novariance mode when it zeroes them out (DiSProD-NV). We look into the state distributions produced by DiSProD and DiSProD-NV, and compare them against empirical state distributions. We fix the start state and a sequence of action distributions, and compute the next-state distribution for a fixed depth. For the empirical state distribution, we sample actions from the same fixed action distribution and use the dynamics model to compute the next-state. The trajectory distributions for Pendulum are visualized in Figure 4. We observe that DiSProD gives us a better approximation than DiSProD-NV - the mean is more accurate and while the variance is underestimated, it has a reasonable shape. Additional experiments in the full paper show additional visualization of trajectory distributions, as well as showing that both action and state variance contribute to the improved planning performance.

We next consider run-time comparing DiSProD to the baselines. Evaluating on the basic Gym environments (details in the full paper), DiSProD-NV has roughly the same run time as CEM and MPPI and DiSProD is up to 7 times slower. This gap is expected to increase with more state variables due to the inclusion of partial derivatives in the computation graph.

5 Conclusion

The paper presents DiSProD, a novel approach for planning in continuous stochastic environments. The method is general and overcomes limitations of prior work by using an abstracted representation, using a higher order Taylor approximation, and showing how optimization via gradients can be done over propagation of distributions. Experiments show success across multiple problems, improved handling of stochastic environments, and decreased sensitivity to search depth, reward sparsity, and large action spaces. DiSProD is also shown to be compatible with control in real robotic systems. At present, the key limitations of DiSProD are its computational complexity arising from incorporating the first and second order partials in the computation graph, over which we perform gradient search, its approximation quality requiring non-zero partial derivatives, and the need for a known model which may be resolved using model-based RL. These are left as important questions for future work.

Acknowledgements

We are grateful to members of the VAIL Lab at Indiana University, especially Durgakant Pushp and Lantao Liu, for helping with the robotic experiments. This work was partly supported by NSF under grants 2002393, 2006886, 2047169 and 2246261. Some of the experiments in this paper were run on the Big Red computing system at Indiana University, supported in part by Lilly Endowment, Inc., through its support for the Indiana University Pervasive Technology Institute.