SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Siyuan Zhou $^{1}$ , Yilun Du $^{2}$ , Shun Zhang $^{3}$ , Mengdi Xu $^{4}$ , Yikang Shen $^{3}$ , Wei Xiao $^{2}$ , Dit-Yan Yeung $^{1}$ , Chuang Gan $^{3,5}$
$^{1}$ Hong Kong University of Science and Technology  
 $^{2}$ Massachusetts Institute of Technology,  $^{3}$ MIT-IBM Watson AI Lab  
 $^{4}$ Carnegie Mellon University,  $^{5}$ UMass Amherst
Diffusion models have risen as a promising approach to data-driven planning, and have demonstrated impressive robotic control, reinforcement learning, and video planning performance. Given an effective planner, an important question to consider is replanning – when given plans should be regenerated due to both action execution error and external environment changes. Direct plan execution, without replanning, is problematic as errors from individual actions rapidly accumulate and environments are partially observable and stochastic. Simultaneously, replanning at each timestep incurs a substantial computational cost, and may prevent successful task execution, as different generated plans prevent consistent progress to any particular goal. In this paper, we explore how we may effectively replan with diffusion models. We propose a principled approach to determine when to replan, based on the diffusion model's estimated likelihood of existing generated plans. We further present an approach to replan existing trajectories to ensure that new plans follow the same goal state as the original trajectory, which may efficiently bootstrap off previously generated plans. We illustrate how a combination of our proposed additions significantly improves the performance of diffusion planners leading to  $38\%$  gains over past diffusion planning approaches on Maze2D, and further enables the handling of stochastic and long-horizon robotic control tasks. Videos can be found on the anonymous website: https://vis-ww.cs.umass.edu/replandiffuser/.
Planning is a pivotal component in effective decision-making, facilitating tasks such as collision avoidance in autonomous vehicles [29], efficient and long-horizon object manipulation [12] and strategic reasoning for extended periods [32, 25]. Recently, diffusion has risen as a promising approach to data-driven planning [17, 1, 14, 40], enabling impressive advancements in robot control [6] and the generation of video-level plans [9].
In many planning settings, an important variable to consider is not only how to plan, but also when it is time to replan. Environments are often stochastic in nature, making previously constructed plans un-executable given an unexpected environmental change. Moreover, inaccurate execution of individual actions in an environment will often also cause states to deviate from those in the expected plan, further necessitating replanning. Finally, environments will often be partially observable, and the optimal action or sequence of actions may deviate substantially from those originally planned, especially when new information comes to light.
Previous works in planning with diffusion models have not explored how to tackle these factors, and typically generate new plans at each iteration. Such a process is computationally expensive as a
(a) Initial Plan
(b) Fast Replanning
Figure 1: Online planning procedure. (a) The agent first generates a plan from scratch at the start of the task. (b) The agent finds the first door inaccessible and replans a new trajectory that deviates slightly from the original plan. (c) The agent finds the yellow door also locked, and replans a completely different trajectory from scratch. generated plan in diffusion models requires hundreds of sampling steps. Frequent replanning may also be detrimental to task completion, as rapid switching between plans impedes persistent progress towards a single final goal [18]. In this paper, we propose a principled approach to adaptively online Replanning with Diffusion Models, called RDM. In RDM, we propose to leverage the internal estimated likelihood of the currently executed plan as a gauge on when is a good time to replan. At the start of plan execution, a predicted plan is freshly sampled from the diffusion model and has a high estimated likelihood. As the plan gets executed, or when unexpected state changes happen, the underlying transitions expected in the plan get increasingly less likely, and the internally estimated likelihood of the diffusion model tells us to replan.
(c) Complete Replanning
Given a previously generated plan, how do we effectively regenerate a new plan? Directly resampling the entire plan can be computationally expensive, as all the previous computations put in the previous plan would be wasted. Furthermore, a newly generated plan might erase all the progress we made previously toward the goal from a previous plan. Instead, based on the feasibility of the current plan being executed, which is estimated by likelihood determined by the diffusion model, we propose to partially replan based on our plan, or fully replan from scratch. We initialize our new plan with the old one and add a small amount of noise to the plan before then denoising the added noise. This light amount of denoising serves as a planning refinement operator.
We illustrate our online planning procedure in Figure 1, where a learned agent first generates a full plan for the goal as in Figure 1a. When the agent reaches the first door, the likelihood of the plan drops, as the agent finds it locked, and the corresponding solution becomes infeasible due to the door's inaccessibility. However, the agent can slightly adjust the initial plan by instead entering the left yellow door and taking a short detour, after which the rest of the plan remains the same. Based on the estimated likelihood of the plan, we may replan and regenerate a new solution based on the previous solution as shown in Figure 1b. In another situation where the yellow door is also locked, and there is no nearby solution, our likelihood plummets, telling us that we need to regenerate a new and completely different solution from scratch as shown in Figure 1c.
In summary, our contributions are three-fold. (1) We introduce RDM, a principled approach to determining when to replan with diffusion models. (2) We illustrate how to effectively replan, using a previously generated plan as a reference. (3) We empirically validate that RDM obtains strong empirical performance across a set of different benchmarks, providing a  $38\%$  boost over past diffusion planning approaches on Maze2D and further enabling effective handling of stochastic and long-horizon robotic control tasks.
2.1 Problem Setting
We consider a trajectory optimization problem similar to Diffuser [17]. A trajectory is represented as  $\tau = (s_0, a_0, s_1, a_1, \ldots, s_T, a_T)$ , where  $s_t, a_t$  are the state and action at the  $t$ -th time step, and  $T$  is the episode length. The dynamics function of the environment is formulated as  $s_{t+1} = f_{\sigma}(s_t, a_t)$  where  $\sigma$  represents the stochastic dynamics. For domains with partial observations, the agent observes an observation  $o_t$  instead of the state  $s_t$ . The objective of the trajectory optimization problem is to find a sequence of actions  $a_{0:T}$  that maximizes  $\mathcal{J}(s_0, a_{0:T})$ , which is the expected sum of rewards
Figure 2: Results overview. The performances of RDM significantly outperform the models with single-step execution (i.e. IQL) and previous planning-based models with diffusion models (i.e. Decision Diffuser) across long-horizon planning (Maze2D), stochastic environments (D4RL locomotion) and robotic control (RLBench). The performance metrics are normalized returns for D4RL tasks (including Maze2D and Locomotion) and success rates for RLBench.
over all the time steps,
$$
\operatorname * {a r g   m a x} _ {a _ {0: T}} \mathcal {J} (s _ {0}, a _ {0: T}) = \operatorname * {a r g   m a x} _ {a _ {0: T}} \operatorname * {E} _ {s _ {1: T}} \sum_ {t = 0} ^ {T} r (s _ {t}, a _ {t}).
$$
Following the Decision Diffuser [1], we plan across state-space trajectories and infer actions through inverse dynamics.
2.2 Planning with Diffusion Models
In Diffuser [17], plan generation is formulated as trajectory generation  $\tau$  through a learned iterative denoising diffusion process  $p_{\theta}(\tau^{i - 1}\mid \tau^i)$  [33]. This learned denoising process is trained to reverse a forward diffusion process  $q(\tau^i\mid \tau^{i - 1})$  that slowly corrupts the structure of trajectories by adding noise. Both processes are modeled as Markov transition probabilities:
$$
q \left(\tau^ {0: N}\right) = q \left(\tau^ {0}\right) \prod_ {i = 1} ^ {N} q \left(\tau^ {i} \mid \tau^ {i - 1}\right), \quad p _ {\theta} \left(\tau^ {0: N}\right) = p \left(\tau^ {N}\right) \prod_ {i = 1} ^ {N} p _ {\theta} \left(\tau^ {i - 1} \mid \tau^ {i}\right). \tag {1}
$$
Given a set of predetermined increasing noise coefficients  $\beta_{1},\ldots ,\beta_{N}$ , the data distribution induced is given by:
$$
q \left(\tau^ {i} \mid \tau^ {i - 1}\right) := \mathcal {N} \left(\sqrt {1 - \beta_ {i}} \tau^ {i - 1}, \beta_ {i} I\right), \quad p _ {\theta} \left(\tau^ {i - 1} \mid \tau^ {i}\right) := \mathcal {N} \left(\mu_ {\theta} \left(\tau^ {i}, i\right), \sigma_ {i} ^ {2} I\right), \tag {2}
$$
where  $q(\tau^0)$  is the data distribution and  $p(\tau^N)$  is a standard Gaussian prior.
Diffuser is trained to maximize the likelihood of trajectories through a variational lower bound over the individual steps of denoising
$$
\mathrm {E} _ {q \left(\tau^ {0}\right)} \left[ \log p _ {\theta} \left(\tau^ {0}\right) \right] \geq \mathrm {E} _ {q \left(\tau^ {0: N}\right)} \left[ \log \frac {p _ {\theta} \left(\tau^ {0 : N}\right)}{q \left(\tau^ {1 : N} \mid \tau^ {0}\right)} \right] \approx \mathrm {E} _ {\tau^ {1: N} \sim q \left(\tau^ {1: N}\right)} \left[ \sum_ {i = 1} ^ {N} \log p _ {\theta} \left(\tau^ {i - 1} \mid \tau^ {i}\right) \right], \tag {3}
$$
where trajectory plans may then be generated by sampling iteratively from the learned denoising network for  $N$  steps. We follow the formulation of planning in Decision Diffuser [1], where if we seek to optimize an objective  $\mathcal{J}(\tau)$ , we directly learn and sampling from conditional trajectory density  $p_{\theta}(\tau|\mathcal{J})$ .
In terms of notations, we follow Janner et al. [17] and use superscripts  $i \in [0, N]$  to denote diffusion steps (for example,  $\tau^i$ ), and use subscripts  $t \in [0, T]$  to denote time steps of the agent's execution in the environment (for example,  $s_t, a_t$ , etc.).
3 Adaptive Online Replanning with Diffusion Models
In this section, we present RDM, an approach to online replanning with diffusion models. First, we discuss different criteria for determining when a solution is problematic and needs replanning. We then introduce the estimated likelihood of a planned trajectory as a criterion to determine when to replan in Section 3.1. In Section 3.2, we discuss how we replan by comparing three distinct replanning strategies and propose an adaptive replanning approach, based on combinations of the two approaches.
Figure 3: Overview of our replanning approach RDM. The first figure shows the likelihood curve of the sampled trajectory as the environmental steps increase. There are two steps  $t$  and  $k$  with low probabilities. Step  $t$  corresponds to Replan from scratch in the middle figure, where RDM regenerates a completely new trajectory based only on the current state and from Gaussian noise. Step  $k$  corresponds to Replan with future in the right-most figure, where RDM shifts the current state as the first state and repeats the last state to fill the rest. RDM also introduces some timesteps of noise to the trajectory and denoises.
3.1 When to Replan
When the agent sequentially executes a plan  $\tau = (s_0, a_0, s_1, a_1, \ldots, s_T, a_T)$ , different factors in the environment such as unexpected interventions or imperfect transitions may cause future states and actions in the plan to deviate substantially from the desired trajectory. We denote with  $\tau_{\rightarrow k}$  the trajectory executed in the environment up to time step  $k$ . Formally,  $\tau_{\rightarrow k} = (s_0, a_0, \bar{s}_1, \bar{a}_1, \ldots, \bar{s}_k, \bar{a}_k, s_{k+1}, \ldots, s_T, a_T)$ , where  $s_t, a_t$  are planned states and actions generated by the diffusion model;  $\bar{s}_t$  is a state observed in the environment, and  $\bar{a}_t$  is an action that is computed using the inverse dynamics during the execution process (which is the action transiting to  $s_{t+1}$  from  $\bar{s}_t$ ). Continuing to execute the current plan will result in more significant discrepancies and potentially catastrophic failures. It may not even be feasible to transit to the next planned state ( $s_{k+1}$ ) from the current state that the agent is at ( $\bar{s}_k$ ). In such cases, we need to update the planned trajectory based on  $\tau_{\rightarrow k}$ , a process we refer to as replanning.
Unfortunately, determining the need for replanning is non-trivial. Directly replanning at a fixed interval is straightforward but inflexible, and may come too late when the environment suddenly changes. It can also make the planning algorithm computationally inefficient when the planned trajectory is executed perfectly and does not require replanning. Instead, it is essential to establish a criterion that automatically determines when replanning is necessary. An alternative approach is to replan when the distance between the actual state observed from the environment and the desired state in the plan exceeds a certain threshold. However, in many cases, even if the actual state the agent is in matches with the state in the plan, new information about the environment may make the current plan infeasible. Thus, it is desirable to come up with a comprehensive and adaptive approach to determine the appropriate timing for replanning.
In this paper, we propose a novel method that utilizes the inherent estimated likelihood of the trajectory in diffusion models as a criterion for replanning. As the diffusion model is trained on different successful trajectories in an environment, its learned likelihood function implicitly captures the executability of an underlying plan. A feasible trajectory will be assigned a high probability, with increasingly less feasible trajectories assigned decreasing probabilities. Thus, when the likelihood assessment of a diffusion model on the plan has a low probability, it is a good indicator of the need for replanning.
The key challenge is to estimate the likelihood of a partially-executed trajectory  $\tau_{\rightarrow k}$ . Recall that we optimize the variational lower bound defined in Equation 3. We propose to use the same objective to determine the likelihood of a partially-executed trajectory. Intuitively, we assess the degree to which a diffusion model recovers its original plan after the plan has been corrupted by Gaussian noise. Concretely, given  $\tau_{\rightarrow k}$ , we added different magnitudes of Gaussian noise to the trajectory, corresponding to sampling based on the diffusion process  $q(\tau^i \mid \tau_{\rightarrow k})$  at a diffusion step  $i$ . We then compute the KL divergence between the posterior  $q(\tau^{i-1} \mid \tau^i, \tau_{\rightarrow k})$  and the denoising step  $p_\theta(\tau^{i-1} \mid \tau^i)$  at the diffusion step  $i$ . We estimate likelihood by computing this KL divergence at several fixed different diffusion steps  $i$ , which we discuss in detail in the appendix.
Algorithm 1 When to Replan.
1: Input: Diffusion model  $p_{\theta}$ , Planning horizon  $H$ , Episode length  $T$ , Threshold  $l_{s}$  of Replanning from scratch, Threshold  $l_{f}$  of Replanning with future and  $l_{s} < l_{f}$ , Diffusion step  $N_{s}$  of Replan from scratch, Diffusion step  $N_{f}$  of Replan with future and  $N_{f} \ll N_{s}$ , Diffusion steps  $I$  for computing likelihood  
2: Environment step  $t = 0$   
3: Generate a trajectory  $\tau_{t} = (s_{t}, a_{t}, \dots, s_{t + H}, a_{t + H})$   
4: for  $t$  in  $0, \dots, T$  do  
5: Extract  $(s_{t}, s_{t + 1}, a_{t})$  from  $\tau_{t}$   
6: Execute  $a_{t}$  and obtain  $s_{t + 1}$  from the environment  
7: Replace the  $s_{t + 1}$  in  $\tau_{t}$   
8:  $\tau^{0} \gets \tau_{t}$   
9: Sample  $\tau^{i} \sim q(\tau^{i} \mid \tau^{0})$  for all  $i \in I$   
10: Compute average likelihood  $L_{t} = \frac{1}{|I|} \sum_{i \in I} D_{KL}(q(\tau^{i - 1} \mid \tau^{i}, \tau^{0}) \| p_{\theta}(\tau^{i - 1} \mid \tau^{i}))$   
11: if  $(t + 1)\%H = 0$  or  $L_{t} <= l_{1}$  then  
12: Sample a trajectory  $\tau_{t}$  with Replanning from scratch  
13: else if  $L_{t} <= l_{2}$  then  
14: Sample a new trajectory  $\tau_{t}$  based on previous trajectory  $\tau_{t - 1}$  with Replanning with future  
15: else  
16:  $\tau_{t + 1} \gets \tau_{t}$   
17: end if  
18: end for
Algorithm 2 Replanning from scratch
1: Input: Diffusion model  $p_{\theta}$ , Diffusion steps  $N_{s}$ , Conditions  $C$   
2: Initialize plan  $\tau^{N_s} \sim \mathcal{N}(\mathbf{0}, \mathbf{I})$   
3: for  $i$  in  $N_{s}, \dots, 1$  do  
4:  $\tau^{i-1} \gets p_{\theta}(\tau^i)$   
5: Constraint  $\tau^{i-1}$  with conditions  $C$   
6: end for  
7: Return:  $\tau^0$   
1: Input: Diffusion model  $p_{\theta}$ , Diffusion process  $q$ , Diffusion steps  $N_f$ , Conditions  $C$ , Previous sampled trajectory  $\tilde{\tau}$   
2:  $\tau^{N_f} \gets q_{N_f}(\tilde{\tau})$   
3: for  $i$  in  $N_f, \dots, 1$  do  
4:  $\tau^{i-1} \gets p_{\theta}(\tau^i)$   
5: Constraint  $\tau^{i-1}$  with conditions  $C$   
6: end for  
7: Return:  $\tau^0$
Given this likelihood estimate, we set a threshold for the likelihood of the planning trajectories. If the likelihood drops below the threshold, it indicates the current plan deviates significantly from the optimal plan and can cause the failure of the plan, which suggests that replanning is necessary. We illustrate this replanning process in Algorithm 1.
3.2 How to Replan
Given a previous partially executed plan executed to time step  $k$ ,  $\tau_{\rightarrow k}$ , and suppose we decide to replan at this current time step using one of the criteria defined the Section 3.1, how do we construct a new plan  $\tau'$ ? We consider three separate replanning strategies, which are the following:
1. Replan from scratch. In this strategy, we regenerate a completely new planning trajectory based only on the current state. This involves completely resampling the trajectory starting from Gaussian noise and running all steps (defined as  $N_{s}$ ) of iterative denoising. Although such an approach fully takes advantage of the diffusion model to generate the most likely trajectory, it is computationally expensive and discards all information in the previously generated plan (Algorithm 2).
2. Replan from previous context. Instead, we may replan by considering both the current and historical states when executing the previously sampled plan. Given a partially-executed trajectory  $\tau_{\rightarrow k}$  and current time step  $k$ , we sample a new trajectory  $\tau'$  based on  $\tau_{\rightarrow k}$ . Specifically, we run  $N_p$  diffusion steps to add noise to  $\tau_{\rightarrow k}$  at all the timesteps that are larger than  $k$  by sampling from the forward process  $q(\tau^{N_p} \mid \tau_{\rightarrow k})$ . We then run  $N_p$  steps of denoising. Here,  $N_p$  is much smaller than Replan from scratch. We will fix the timesteps of  $\tau_{\rightarrow k}$  that are smaller than  $k$  to their observed
value, as we have already executed them in the environment. This approach can be seen as locally repairing a predicted plan, and thus requires a substantially smaller number of diffusion steps.
3. Replan with future context. In a similar way to replan with previous context, in replan with future context, we initialize our new plan  $\tau'$  based on our partially-executed plan  $\tau_{\rightarrow k}$ . However, in contrast to the previous approach, we do not add the previously executed states of the plan. We remove the executed states in  $\tau_{\rightarrow k}$  and repeatedly add the last predicted state in  $\tau_{\rightarrow k}$  on the right to fill in the remaining plan context. Namely,  $\tilde{\tau} = (\bar{s}_k, \bar{a}_k, s_{k+1}, a_{k+1}, \dots, s_T, a_T, s_T, a_T, \dots, s_T, a_T)$ , where  $s_T, a_T$  are repeated to match the input length of the diffusion model. To replan, similar to replan from previous context, we add  $N_f$  steps of noise to  $\tilde{\tau}$  by sampling from the forward process  $q(\tau^{N_f} \mid \tau_{\rightarrow k})$  and then denoise for  $N_f$  steps (Algorithm 3), where  $N_f$  is much smaller than  $N_s$ .
Replanning from scratch is highly time-consuming and inefficient for real-time tasks since it requires all diffusion steps to be run to generate a plan from Gaussian noise. However, in cases of significant interventions, it becomes necessary to regenerate the plan from scratch. Both replanning for future context or past context are efficient and locally repair the previously generated plan. Replanning with past context remains more consistent with the original goal of the previous plan, while replanning for future context further extends the current plan to future context.
We empirically find that replanning with future context is a more effective approach. In our proposed approach, based on the estimated likelihood of a current plan, we either replan from scratch or replan with future context. We illustrate the psuedocode for our approach in Algorithm 1. We provide a detailed ablation analysis of the different replanning strategies in the experiment section (Figure 7).
In this section, we empirically use our proposed RDM algorithm for replanning in multiple decision-making tasks (as illustrated in Figure 2) and evaluate the performance of the generated trajectories. We empirically answer the following questions: Q1: In domains with long horizon plans, does RDM effectively and efficiently fix the initial plan? Q2: In domains with unexpected interventions, which makes the initial plan no longer feasible during the execution, does RDM efficiently find a feasible plan? Q3: How do when to replan and different replanning strategies affect the agent's performance? In all the experiments below, we report the results using five random seeds.
4.1 Long Horizon Planning
We first evaluate the long-horizon planning capacity of RDM and the baseline methods in the Maze2D environment. In this environment, the large Maze2D domain is especially challenging, as it may require hundreds of steps to reach the target and we found that baseline algorithms struggled to find perfect trajectories given the long horizon and sparse reward setting.
Maze2D is a goal-conditioned planning task, where given the starting location and the target location, the agent is expected to find a feasible trajectory that reaches the target from the starting location avoiding all the obstacles. The agent receives a reward of 1 when it reaches the target state and receives a reward of 0 otherwise. The starting and target locations are randomly generated, consistent with the setting in the Diffuser [17]. We evaluate replanning using our RDM algorithm and multiple competing baseline algorithms, including model-free reinforcement learning algorithms such as Batch-Constrained deep Q-learning (BCQ) [11], Constrained Q-learning (CQL) [22], Implicit Q-learning (IQL) [21], and planning-based algorithms like Diffuser [17].