SlowGuess/NeurIPS_2023_OCR huggingface dataset:
We report the results in Table 1 where numbers for all baselines are taken from Diffuser [17]. As expected, model-free reinforcement learning algorithms (BCQ, CQL, and IQL) struggle to solve long-horizon planning tasks. On the other hand, Diffuser generates a full plan and outperforms these approaches, but does not replan online. As a result, it still does not perform well when the size of the domain increases. As the planning horizon increases, it is likely that some of the transitions generated by the diffusion model are infeasible as shown in Figure 4c. We find that our replanning approach, RDM, significantly outperforms all the baseline algorithms. Specifically, it achieves a score of  $190\%$  in the Maze2D Large setting, which is almost  $63\%$  higher than the Diffuser. We also visualize the sampled trajectory and the executed trajectories in Figure 4. We can observe that infeasible transitions and collisions with the wall will drop the likelihood substantially. RDM is able to find these infeasible transitions and replan accordingly.
(a) Likelihood
(b) Sample
Figure 4: Visualization on Maze2D. (a) illustrates the likelihood curve of the sampled trajectory as the environmental steps increase, with a noticeable drop when the agent collides with the wall. (b) presents the sampled trajectory at the onset of the task. (c) demonstrates the actual trajectory of Diffuser, which follows a fixed interval replanning strategy. It is evident that the agent does not recognize the flawed part of the sampled trajectory, leading it to a less favorable state and resulting in wasted environmental steps. (d) shows the actual trajectory of RDM. In contrast to (c), the agent successfully detects collisions and replans a feasible trajectory.
(c) Diffuser
(d) RDM
Environment	BCQ	CQL	IQL	Diffuser	RDM (Ours)
Maze2D	U-Maze	12.8	5.7	47.4	113.9	133.6 ±3.0
Maze2D	Medium	8.3	5.0	34.9	121.5	159.3 ±3.1
Maze2D	Large	6.2	12.5	58.6	123.0	185.3 ±11.7
Average	9.1	7.7	47.0	119.5	159.4
Multi2D	U-Maze	-	-	24.8	128.9	152.1 ±2.31
Multi2D	Medium	-	-	12.1	127.2	155.1 ±4.0
Multi2D	Large	-	-	13.9	132.1	195.4 ±7.6
Average	-	-	16.9	129.4	167.5
Table 1: Long horizon planning. The normalized returns of RDM and other baselines on Maze2D. RDM achieves  $38\%$  gains beyond Diffuser for all tasks and even  $63\%$  gains for large maps.
4.2 Stochastic Environments
Next, we investigate whether RDM can tackle environments with stochastic dynamics. We consider a stochastic variation of the D4RL locomotion benchmark [10] by adding randomness to the transition function. Stochasticity can make the trajectory optimization problem more challenging, as the executed trajectory can be significantly different from the initially planned trajectory. We omitted the Half cheetah environment as we found all algorithms performed poorly with added stochasticity. In this setting, we consider competing baseline methods including a state-of-the-art model-free algorithm, IQL [21], a Transformer-based RL algorithm, Decision Transformer (DT) [5], and a diffusion model-based algorithm, Decision Diffuser (DD) [1].
As shown in Table 2, although the performance of all the methods drops significantly compared with one without stochasticity, RDM performs comparatively well among all baselines. RDM with replanning strategy is able to detect such stochasticity and quickly regenerate a new plan, resulting in improved performance. We compare DD and RDM using the same number of total diffusion resampling steps.
4.3 Robotic Control
We further evaluate the RDM's capacity in complex domains with visual observations. We consider the RLBench domain [15], which consists of 74 challenging vision-based robotic learning tasks. The agent's observations are RGB images captured by multi-view cameras. The agent controls a 7 DoF robotic arm. The domain is designed to reflect a real-world-like environment with high dimensionality in both the observation space and the action space, making it challenging to solve (Figure 5b).
Most previous works tackle the RLBench tasks by taking actions in macro-steps [13], where the agent switches to a different macro-step when the state of the gripper or the velocity changes. However, such an approach has several limitations: (1) The longest trajectory length achievable using macro steps is relatively short, which is slightly over 30 steps. Such a limitation hinders the agent's ability
Dataset	Environment	IQL	DT	DD	RDM (Ours)
Med-Expert	Hopper	48.9	52.5	49.4	59.4 ±3.4
Med-Expert	Walker2d	90.7	90.4	71.1	92.5 ±2.2
Medium	Hopper	47.6	49.9	45.3	53.2 ±1.2
Medium	Walker2d	71.8	67.9	60.7	70.7 ±3.0
Med-Replay	Hopper	55.3	40.8	47.2	57.9 ±3.2
Med-Replay	Walker2d	54.4	52.6	46.4	57.3 ±3.1
Average	61.5	59.0	53.4	65.2
Fail to open box Replan and retry
(a) With our replanning approach
Fail to open box Continue the plan
(b) Without our replanning approach
Figure 5: Comparison between trajectories with and without replanning. (a) Our replanning approach allows the agent to detect when the robotic arm fails to open the box and instructs it to retry. (b) Without our replanning approach, the failure goes unnoticed, leading the arm to persist in following the previous, unsuccessful plan.
to solve long-horizon tasks. (2) Macro-steps are pre-defined, which usually require expert knowledge or demonstrations. (3) Macro-steps constrain the agent's action space, which makes it impossible to optimize at a finer granularity (for example, avoiding an intermediate collision) and prevent the agent from doing fine-grained control. Therefore, we do not use macro-step in our experiments.
We choose the following three tasks for evaluation: open box, close box, and close fridge. The experimental results are in Table 3. We find that baseline algorithms fail to accomplish these three tasks. In contrast, RDM outperforms all the baselines by a large margin and achieves a success rate of  $77.3\%$ . We visualize the execution of RDM in Figure 5. RDM successfully detects failures in the initial plan and efficiently replans during the execution.
Table 2: Stochastic environments. The normalized returns of RDM and other baselines on Locomotion with randomness. RDM exhibits comparable or superior performance.
Environment	BC	IQL	DT	DD	RDM (Ours)
open box	9.3	17.6	9.8	52.9	72.0 ±7.5
close box	5.9	10.3	7.8	46.0	61.3 ±7.7
close fridge	13.0	23.1	21.7	84.6	98.6 ±2.7
Average	9.4	17.0	13.1	61.2	77.3
Table 3: Robotic control. The success rate of RDM and other baselines for three tasks: open box, close box and close fridge on RLBench. RDM outperforms baselines by a large margin.
4.4 Ablative Study
Different Levels of Stochasticity. We investigate how the algorithms respond to different levels of stochasticity by setting different noise values  $\epsilon$  in the task Hopper-medium-expert, where at each time step, the transition is the same as the deterministic version of the domain with probability  $1 - \epsilon$ ,
Figure 6: Different levels of stochasticity. The normalized returns with different probabilities  $\epsilon$  of randomness in the task Hopper-medium-expert. By adaptive planning, RDM is more effective than other baselines.
Figure 7: Replan analysis. We analyze replanning at fixed intervals (DD), replanning based on state distance (SDM), and replanning based on past context (RPDM). With the same number of diffusion steps on Hopper-medium-expert, RDM obtains the best performance.
and is random (as if a random action is taken) with probability  $\epsilon$ . The results are in Figure 6. We find that RDM consistently outperforms other baselines under all  $\epsilon$  values.
Different Replanning Strategies. To analyze the different replan approaches, we consider the following baselines. (1) Decision Diffuser (DD) replans using two fixed intervals, for deciding when to replan from scratch and when to plan with future context, respectively. (2) Replan on State Distance (SDM) replans when the error between the observed states  $(\bar{s}_t)$  and desired states in the plan  $(s_t)$  exceeds a threshold. It replans with the future context. The following two algorithms determine when to plan according to the diffusion model likelihood. (3) Replan from previous context (RPDM) replans with on past context, and (4) RDM replans with future context. In Figure 7, we find that SDM and RPDM both have a better performance than DD, among which RDM achieves the best performance. All methods use the same number of total diffusion steps.
We provide additional ablation studies and visualization analysis in the Appendix.
Diffusion Models for Decision Making. The use of diffusion models for decision-making tasks has seen increased interest in recent years [17, 1, 9, 6, 30, 40, 23, 14, 38, 24] and have been applied to tasks such as planning [17, 1], robotic control [6], behavior cloning [30], and video generation [9]. In [17, 1] introduce an approach to planning with diffusion models, but both methods typically replan at each iteration. In this paper, we propose a method to adaptively determine when to replan with diffusion models and illustrate different approaches to replan.
Model-Based Planning. Our work is also related to existing work in model-based planning [35, 4, 7, 19, 3, 36, 2, 27, 34, 8, 37, 28]. For domains with finite states, traditional methods use dynamic programming to find an optimal plan, or use Dyna-algorithms to learn the model first and then generate a plan [35]. For more complex domains, a learned dynamics model be autoregressively rolled out [7, 19], though gradient planning through these models often lead to adversarial examples [3, 36, 2]. To avoid this issue approaches have tried gradient-free planning methods such as random shooting [27], beam search [34], and MCMC sampling [4]. Our work builds on a recent trend in model-based planning which folds the planning process into the generative model [16].
6 Limitations and Conclusion
Conclusion. In this paper, we propose RDM, an adaptive online replanning method, with a principled approach to determine when to replan and an adaptive combined strategy for how to replan. We demonstrate RDM's remarkable proficiency in managing long-horizon planning, unpredictable dynamics, and intricate robotic control tasks. An interesting line of future work involves investigating the scalability of RDM to more complex and larger-scale problems and extending to real robot tasks and autonomous driving.
Limitations. There are several limitations to our approach. Firstly, the effectiveness of planning and replanning is heavily dependent on the generative capabilities of the diffusion models. If the diffusion
model is unable to accurately represent the current task, its plans may not be viable and RDM may fail to identify problematic trajectories. Secondly, despite RDM's significant improvement on the performance of previous diffusion planning methods, we observed that likelihood-based replanning was insensitive to some planning failures. This highlights the potential benefit of combining likelihood-based replanning with other indicators for replanning.
This work has been made possible by a Research Impact Fund project (R6003-21) funded by the Research Grants Council of the Hong Kong Government.
[1] A. Ajay, Y. Du, A. Gupta, J. Tenenbaum, T. Jaakkola, and P. Agrawal. Is conditional generative modeling all you need for decision-making? arXiv preprint arXiv:2211.15657, 2022.
[2] K. Asadi, D. Misra, and M. Littman. Lipschitz continuity in model-based reinforcement learning. In International Conference on Machine Learning, 2018.
[3] S. Bengio, O. Vinyals, N. Jaitly, and N. Shazeer. Scheduled sampling for sequence prediction with recurrent neural networks. In Advances in Neural Information Processing Systems, 2015.
[4] H. Chen, Y. Du, Y. Chen, J. Tenenbaum, and P. A. Vela. Planning with sequence models through iterative energy minimization. arXiv preprint arXiv:2303.16189, 2023.
[5] L. Chen, K. Lu, A. Rajeswaran, K. Lee, A. Grover, M. Laskin, P. Abbeel, A. Srinivas, and I. Mordatch. Decision transformer: Reinforcement learning via sequence modeling. Advances in neural information processing systems, 34:15084-15097, 2021.
[6] C. Chi, S. Feng, Y. Du, Z. Xu, E. Cousineau, B. Burchfiel, and S. Song. Diffusion policy: Visuomotor policy learning via action diffusion. arXiv preprint arXiv:2303.04137, 2023.
[7] S. Chiappa, S. Racaniere, D. Wierstra, and S. Mohamed. Recurrent environment simulators. arXiv preprint arXiv:1704.02254, 2017.
[8] Y. Du, T. Lin, and I. Mordatch. Model based planning with energy based models. In Conference on Robot Learning, 2019.
[9] Y. Du, M. Yang, B. Dai, H. Dai, O. Nachum, J. B. Tenenbaum, D. Schuurmans, and P. Abbeel. Learning universal policies via text-guided video generation. arXiv e-prints, pages arXiv-2302, 2023.
[10] J. Fu, A. Kumar, O. Nachum, G. Tucker, and S. Levine. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
[11] S. Fujimoto, D. Meger, and D. Precup. Off-policy deep reinforcement learning without exploration. In International conference on machine learning, pages 2052-2062. PMLR, 2019.
[12] C. R. Garrett, R. Chitnis, R. Holladay, B. Kim, T. Silver, L. P. Kaelbling, and T. Lozano-Pérez. Integrated task and motion planning. Annual review of control, robotics, and autonomous systems, 4:265-293, 2021.
[13] P.-L. Guhur, S. Chen, R. G. Pinel, M. Tapaswi, I. Laptev, and C. Schmid. Instruction-driven history-aware policies for robotic manipulations. In Conference on Robot Learning, pages 175-187. PMLR, 2023.
[14] S. Huang, Z. Wang, P. Li, B. Jia, T. Liu, Y. Zhu, W. Liang, and S.-C. Zhu. Diffusion-based generation, optimization, and planning in 3d scenes. arXiv preprint arXiv:2301.06015, 2023.
[15] S. James, Z. Ma, D. R. Arrojo, and A. J. Davison. Rlbench: The robot learning benchmark & learning environment. IEEE Robotics and Automation Letters, 5(2):3019-3026, 2020.
[35] R. S. Sutton and A. G. Barto. Reinforcement learning: An introduction. MIT press, 2018.
[36] E. Talvitie. Model regularization for stable sample rollouts. In UAI, pages 780-789, 2014.
[37] A. Tamar, Y. Wu, G. Thomas, S. Levine, and P. Abbeel. Value iteration networks. Advances in neural information processing systems, 29, 2016.
[38] Z. Wang, J. J. Hunt, and M. Zhou. Diffusion policies as an expressive policy class for offline reinforcement learning. arXiv preprint arXiv:2208.06193, 2022.
[39] Y. Wu and K. He. Group normalization. In European Conference on Computer Vision, 2018.
[40] E. Zhang, Y. Lu, W. Wang, and A. Zhang. Lad: Language augmented diffusion for reinforcement learning. arXiv preprint arXiv:2210.15629, 2022.
(a) Problematic Sample
(b)  $N = 40$
Figure 8: Samples of replanning with different diffusion steps. (a) demonstrates the problematic trajectory. (b) presents the sampled trajectory after replanning with  $N = 40$ . The sampled trajectory still leads to the collision. (c) presents the sampled trajectory after replanning with  $N = 80$ . The trajectory becomes feasible and avoids collision. (d) shows the sampled trajectory after replanning with  $N = 256$ . The trajectory improves quality but needs much more time.
(c)  $N = 80$
(d)  $N = 256$
In the supplementary, we first discuss the experimental details and hyperparameters in Section A. Next, we analyze the impact of different numbers of diffusion steps  $N$  on the replanning process in Section B, and further present the visualization in RLBench in Section C. Finally, we discuss how to compute the likelihood in Section D.
1. Our architecture is built based on Diffuser [17] and Decision Diffuser [1]. In detail, our architecture comprises a temporal U-Net structure with six repeated residual networks. Each network consists of two temporal convolutions followed by GroupNorm [39], and a final Mish nonlinearity [26]. Additionally, We incorporate timestep and conditions embeddings, which are both 128-dimensional vectors produced by MLP, within each block.
2. In RLBench, we exploit the image encoder from the pre-trained CLIP [31], followed by a 2-layered MLP with 512 hidden units and Mish activations.
3. The model is trained using Adam optimizer [20] with a learning rate of 2e-04 and a batch size of 64 for 1e6 training steps.
4. The planning horizon is set to 128 in Maze2D//Multi2D U-Maze, 256 in Maze2D//Multi2D Medium, 256 in Maze2D//Multi2D Large, 64 in Stochastic Environments, and 64 in RLBench.
5. We use a threshold of 0.7 for Replan from scratch and a threshold of 0.5 for Replan with future.
6. The probability  $\epsilon$  of random actions is set to 0.03 in Stochastic Environments.
7. The diffusion steps  $i$  for computing likelihood is set to  $\{5, 10, 15\}$  in Maze2D and Stochastic Environments and  $\{10, 20, 30\}$  in RLBench.
8. The total number of diffusion steps, corresponding to the number of diffusion steps for Replan from scratch is set to 256 in Maze2D, 200 in Stochastic Environments, and 400 in RLBench. And the number of diffusion steps for Replan with future is set to 80 in Maze2D, 40 in Stochastic Environments, and 100 in RLBench.
9. We perform the whole experiment with a total of three Tesla V100 GPUs.
Environment	Diffuser	RDM (N=40)	RDM (Ours)	RDM (N=256)
Multi2D	Large	129.4	160.9	195.4	197.7
Table 4: Replanning with different diffusion steps. RDM with  $N = 40$  performs better than Diffuser. Our method RDM with  $N = 80$  further improves performance and is much more efficient than RDM with  $N = 256$ .
B Different Diffusion Steps for Replanning
In this section, we visualize the impact of using different numbers of diffusion steps  $N$  to replan an existing trajectory. Figure 8 illustrates a problematic sampled trajectory after execution. When
(b)  $N = 40$
Figure 9: Samples of replanning with different diffusion steps. (a) demonstrates the problematic trajectory. (b)(c) presents the sampled trajectory after replanning with  $N = 40$  and 80. The previously sampled trajectory is so troublesome that it can't recover from fast replanning. So the sampled trajectories still lead to the collision. (d) shows the sampled trajectory after replanning with  $N = 256$ . After Replan from scratch, the resampled trajectory becomes feasible.
Fail to close box Replan and retry
(a) With our replanning approach
Fail to close box Continue the plan
(b) Without our replanning approach
Figure 10: Comparison between trajectories with and without replanning. (a) Our replanning approach allows the agent to detect when the robotic arm fails to close the box and instructs it to retry. (b) Without our replanning approach, the failure goes unnoticed, leading the arm to persist in following the previous, unsuccessful plan.
$N = 40$  of replanning, the resampled trajectory still leads to the collision with the wall. However, with  $N = 80$ , which aligns with the number of diffusion steps used in Replan with future, we observe that the sampled trajectory becomes feasible. For  $N = 256$ , corresponding to the number of diffusion steps used in Replan from scratch, the resampled trajectory shows additional improvement, but such a procedure is substantially more expensive than our replanning procedure.
We illustrate another example plan which based off likelihood would need to be planned from scratch in Figure 9. After replanning with  $N = 40$  and  $N = 80$ , the trajectories still encounter issues with colliding with the wall. However, With  $N = 256$ , RDM successfully regenerates a completely different and feasible trajectory through replanning from scratch.
We further evaluate the performance with different replanning steps in Table 4. The results clearly demonstrate that that RDM consistently outperforms Diffuser in all cases. Remarkably, RDM with  $N = 80$ , corresponding to Replan with future, achieves performance close to RDM with  $N = 256$  which corresponds to Replan from scratch which is substantially more expensive.
C Comparison between trajectories with and without replanning
In this section, we visualize the execution in the task close box in Figure 10. In the second frame of Figure 10, we observe that the robotic arm fails to close the box due to insufficient contact between the gripper and the lid. In Figure 10a, RDM successfully detects the failures and retries to establish
(a)  $i = 60$
(b)  $i = 30$
Figure 11: Different diffusion steps  $i$  for computing KL divergence. We estimate the likelihood of trajectory by computing KL divergence between  $q(\tau^{i-1} \mid \tau^i, \tau_{\rightarrow k})$  and  $p_\theta(\tau^{i-1} \mid \tau^i)$ . Here, we illustrate the curve of KL divergence at different environmental steps in Maze2D where the number of total diffusion steps is 256. (a) When large amounts of noise are added at  $i = 60$ , the agent fails to detect any infeasibility as the KL divergence remains low and consistent. (b) When intermediate amounts of noise are added at  $i = 30$ , the agent has the ability to notice infeasibility but the value of KL divergence is similar to that of feasible trajectories. (c) When small amounts of noise are added at  $i = 15$ , the agent can effectively discern infeasibilities, as the KL divergence exhibits a clear boundary between feasible and infeasible trajectories.
(c)  $i = 15$
proper contact with the lid. As a result, the gripper successfully closes the box this time. In contrast, Figure 10b demonstrates the scenario where the arm fails to notice the failures, making the arm persist in the unsuccessful plan. These results demonstrate the capability of RDM to detect failures and replan a successful trajectory in challenging robotic control tasks.
D How to Compute Likelihood
In this section, we analyze the number of noising steps we should use to estimate the likelihood. The likelihood estimation is performed by computing the KL divergence between  $q(\tau^{i - 1} \mid \tau^i, \tau_{\rightarrow k})$  and  $p_{\theta}(\tau^{i - 1} \mid \tau^i)$  with different magnitudes of noise added at diffusion timesteps  $i$ . We analyze the impact of different diffusion steps  $i$  and the corresponding magnitude of noise added for computing KL divergence in Figure 11. We observe that (1) When large amounts of noise are added at  $i = 60$  the agent fails to detect any infeasibility, as the KL divergence remains low and consistent. (2) When intermediate amounts of noise are added at  $i = 30$ , the agent has the ability to notice infeasibilities, but the KL divergence value is similar to that of a feasible trajectory. (3) When small amounts of noise are added at  $i = 15$ , the agent can effectively discern infeasibilities, as the KL divergence exhibits a clear boundary between feasible and infeasible trajectories. Based on these findings, we choose smaller values of  $i$  for computing the likelihood of trajectories and estimating when to replan.
We show the additional experimental results in the following.
1. We show the results of the comparison between our RDM algorithm and the replanning-based baseline algorithms (SDM and PRDM) across more tasks. We run experiments using different intervals or thresholds and measure their computational costs by the total number of diffusion steps. Notably, RDM consistently outperforms all the baseline algorithms under the same total diffusion steps (that is, with the same computational budget).
2. We investigate different intervals  $I$  for replanning for Diffuser and Decision-Diffuser. We observe that as the interval decreases (that is replanning is done more frequently), the performance improves as we expect. However, when the interval is smaller than a certain value, the performance decreases significantly. For example, in the Maze2D Large domain shown in Figure VI, the return increases when the interval decreases from 250 to 100, while it drops when the interval decreases from 100 to 1. The ablation results confirm our statement that replanning with an interval that is too small (for example, replanning at every time step) may prevent successful task execution.
3. We also analyze the impact of different thresholds  $l$  for the baseline algorithms, SDM and RDM. As we expect, when  $l$  decreases (that is, when replanning is done more frequently), the performance of all the methods has improved. Despite the performance differences of the baseline algorithms, we want to emphasize again that under the same computational budget, our RDM algorithm outperforms all the baseline algorithms.