SlowGuess/NeurIPS_2023_OCR huggingface dataset:
To mitigate the issue of adversarial guidance, we present a regularization method that prevents the gap predictor from directing plans in the wrong direction. Inspired by the prior studies which improve the model performance by utilizing attribution maps (Nagisetty et al., 2020; Bertoin et al., 2022), we measure a total variation of the attribution map  $M$  obtained from any input attribution methods  $M = E(\mathcal{G}_{\psi}(\tau ,t))$ . Each element of the attribution map indicates the extent to which the final prediction is influenced by the corresponding input feature. The rationale of employing the total variation of  $M$  lies in the
hypothesis that transitions with excessively high attribution scores are more likely to be outliers. This is because a sequence of transitions within a planned trajectory, rather than a single one, causes a plan to have a high restoration gap. By adding this attribution map regularization, Equation 21 becomes:
$$
\left. \mathrm {d} \boldsymbol {\tau} = \left[ \mathbf {f} (\boldsymbol {\tau}, t) - g (t) ^ {2} \left(\mathbf {s} _ {\theta} (\boldsymbol {\tau}, t) + \alpha \left(\nabla \mathcal {J} _ {\phi} (\boldsymbol {\tau}) - \beta \nabla \mathcal {G} _ {\psi} (\boldsymbol {\tau}, t) - \lambda \nabla \| \nabla M \|\right)\right) \right] \mathrm {d} t + g (t) \mathrm {d} \bar {\mathbf {w}}, \right. \tag {22}
$$
Figure 3: Planning performance of RGG+ on Maze2D-Large single-task with varying  $\lambda$  values.
where  $\lambda$  is a control parameter given by a positive constant, encouraging the attribution map to have a simple, organized structure while preventing the occurrence of adversarial artifacts. We refer to this modification as RGG+.
Table 1: The performance of RGG, RGG+, and various previous algorithms, measured as normalized average return, is presented on the D4RL locomotion benchmark (Fu et al., 2020). Results for RGG and RGG+ show the mean and standard error over 15 planning seeds. Detailed sources for the performance of prior methods are provided in Appendix E.
Dataset	Environment	BC	CQL	IQL	DT	TT	MOPO	MORel	MBOP	Diffuser	RGG	RGG+
Med-Expert	HalfCheetah	55.2	91.6	86.7	86.8	95.0	63.3	53.3	105.9	79.8	90.8 ± 0.3	91.2 ± 0.3
Med-Expert	Hopper	52.5	105.4	91.5	107.6	110.0	23.7	108.7	55.1	107.2	109.6 ± 2.3	109.9 ± 2.3
Med-Expert	Walker2d	107.5	108.8	109.6	108.1	101.9	44.6	95.6	70.2	108.4	107.8 ± 0.1	107.7 ± 0.2
Medium	HalfCheetah	42.6	44.0	47.4	42.6	46.9	42.3	42.1	44.6	44.2	44.0 ± 0.3	44.2 ± 0.3
Medium	Hopper	52.9	58.5	66.3	67.6	61.1	28.0	95.4	48.8	58.5	82.5 ± 4.3	84.9 ± 4.1
Medium	Walker2d	75.3	72.5	78.3	74.0	79.0	17.8	77.8	41.0	79.7	81.7 ± 0.5	82.0 ± 0.4
Med-Replay	HalfCheetah	36.6	45.5	44.2	36.6	41.9	53.1	40.2	42.3	42.2	41.0 ± 0.2	41.3 ± 0.2
Med-Replay	Hopper	18.1	95.0	94.7	82.7	91.5	67.5	93.6	12.4	96.8	95.2 ± 0.5	95.8 ± 0.5
Med-Replay	Walker2d	26.0	77.2	73.9	66.6	82.6	39.0	49.8	9.7	61.2	78.3 ± 4.4	77.5 ± 4.7
Average	51.9	77.6	77.0	74.7	78.9	42.1	72.9	47.8	75.3	81.2	81.6
We present the analytical results of approaches to improve planning performance by leveraging guidance from our proposed metric, the restoration gap, for a wide range of decision-making tasks in offline control settings. Specifically, we demonstrate (1) the relationship between a high restoration gap and poor planning performance, (2) the enhancement of planning performance in the diffusion planner by leveraging restoration gap guidance, and (3) explainability by presenting the attribution maps of the learned gap predictor, highlighting infeasible transitions. More information about our experimental setup and implementation details can be found in Appendix C and Appendix D, respectively.
5.1 Relationship between Restoration Gap and Planning Performance
We evaluate how effectively the restoration gap can identify infeasible plans by comparing our metric with a realism score (Kynkänniemi et al., 2019) and rarity score (Han et al., 2023). Both prior metrics are designed to assess the quality of generated samples by examining the discrepancy between the generated sample and the real data manifold in the feature space. Figure 4 illustrates the performance of plans which are chosen up to top- $\mathrm{K\%}$  from each metric. As illustrated in Figure 4, the higher the restoration gap of the plan is, the poorer the performance is, which implies that the restoration gap captures the quality of the plan well compared to other metrics.
Figure 4: Performance of plans chosen from the top-K % considering various metrics, in the Maze2D-Large single-task.
5.2 Planning Performance Enhancement
Maze2D Experiments Maze2D environments (Fu et al., 2020) involve a navigation task that requires an agent to exhibit long-horizon planning abilities to reach a target goal location. Maze2D environments consist of two tasks: a single-task where the goal location is fixed, and a multi-task which we refer to as Multi2D where the goal location is randomized at the beginning of every episode. We compare our methods
with the model-free offline reinforcement learning algorithms CQL (Kumar et al., 2020) and IQL (Kostrikov et al., 2022); conventional trajectory optimizer MPPI (Williams et al., 2015); and sequence modeling approach Diffuser (Janner et al., 2022). As shown in Table 2, RGG improves the planning performance of Diffuser in 5 out of 6 tasks, with notable improvements in the Maze2D-Large envi
Table 2: Diffuser with RGG and Diffuser with RGG+ outperform all baselines. We report the mean and the standard error over 1000 planning seeds.
Environment	CQL	MPPI	IQL	Diffuser	RGG	RGG+
Maze2D	U-Maze	5.7	33.2	47.4	108.6 ± 1.4	108.8 ± 1.4	109.5 ± 1.3
Maze2D	Medium	5.0	10.2	34.9	129.8 ± 0.7	131.8 ± 0.5	132.1 ± 0.4
Maze2D	Large	12.5	5.1	58.6	123.5 ± 2.0	135.4 ± 1.7	143.9 ± 1.5
Single-task Average	7.7	16.2	47.0	120.6	125.3	128.5
Multi2D	U-Maze	-	41.2	24.8	127.9 ± 0.8	128.3 ± 0.8	128.3 ± 0.8
Multi2D	Medium	-	15.4	12.1	130.1 ± 0.9	130.0 ± 0.9	130.0 ± 0.9
Multi2D	Large	-	8.0	13.9	141.2 ± 1.6	148.3 ± 1.4	150.9 ± 1.3
Multi-task Average	-	21.5	16.9	133.1	135.5	136.4
ronments where the complexity of the obstacle maps is higher than in U-Maze or Medium layouts, leading to a higher occurrence of infeasible plans. RGG+ performs on par with or better than RGG. In contrast, model-free algorithms fail to reliably achieve the goal, as Maze2D environments require hundreds of steps to arrive at the goal location.
Locomotion Experiments Gym-MuJoCo locomotion tasks (Fu et al., 2020) are standard benchmarks in evaluating algorithms on heterogeneous data with varying quality. We compare our methods with the model-free algorithms CQL (Kumar et al., 2020) and IQL (Kostrikov et al., 2022); model-based algorithms MOPO (Yu et al., 2020), MOReL (Kidambi et al., 2020), and MBOP (Argenson & Dulac-Arnold, 2021); sequence modeling approach Decision Transformer (DT) (Chen et al., 2021), Trajectory Transformer (TT) (Janner et al., 2021) and Diffuser (Janner et al., 2022); and pure imitation-based approach behavior-cloning (BC). As indicated in Table 1, our approach of refining Diffuser with RGG either matches or surpasses most of the offline RL baselines when considering the average score across various tasks. Additionally, it significantly enhances the performance of Diffuser, particularly in the "Medium" dataset. We attribute this improvement to the sub-optimal and exploratory nature of the policy that was used to generate the "Medium" dataset, which results in a challenging data distribution to learn the diffusion planner. Consequently, RGG clearly contributes to the enhancement of planning performance. However,  $\mathrm{RGG + }$  only brings about a marginal improvement over RGG. This might be because we adopt the strategy of (Janner et al., 2022), using a closed-loop controller and a shorter planning horizon in locomotion environments compared to Maze2D environments, thereby simplifying the learning process of the gap predictor.
Block Stacking Experiments The block stacking task suite with a Kuka iwa robotic arm is a benchmark to evaluate the model performance for a large state space (Janner et al., 2022) where the offline demonstration data is achieved by PDDLStream (Garrett et al., 2020). It involves two tasks: an unconditional stacking task whose goal is to
Table 3: The performance of RGG, RGG+, and various prior methods evaluated over 100 planning seeds. A score of 100 is desired, while a random approach would receive a score of 0.
Environment	BCQ	CQL	Diffuser	RGG	RGG+
Unconditional Stacking	0.0	24.4	53.3 ± 2.4	63.3 ± 2.7	65.3 ± 2.0
Conditional Stacking	0.0	0.0	44.3 ± 3.2	53.0 ± 3.3	56.7 ± 3.1
Average	0.0	8.1	48.8	58.2	61.0
maximize the height of a block tower, and a conditional stacking task whose goal is to stack towers of blocks subject to a specified order of blocks. We compare our methods with model-free offline reinforcement learning algorithms BCQ (Fujimoto et al., 2019) and CQL (Kumar et al., 2020), and Diffuser (Janner et al., 2022). We present quantitative results in Table 3, where a score of 100 corresponds to the successful completion of the task. The results demonstrate the superior performance of RGG over all baselines, with  $\mathrm{RGG + }$  further enhancing this planning performance.
5.3 Injecting Explainability to Diffusion Planners
The explainability of decision-making models is particularly important in control domains as they could potentially harm physical objects including humans (Kim & Choi, 2021; Lee et al., 2023; Beechey et al., 2023; Kim et al., 2023; Kenny et al., 2023). Training the gap predictor enables the diffusion planner to have explainability. Diffusion planners often generate trajectories with unreliable transitions resulting in execution failures. Attribution maps from the gap predictor highlight such unreliable transitions by identifying the extent to which each transition contributes to the decision of the gap predictor. Specifically, in Maze2D, the attribution maps emphasize the transitions involving wall-crossing or abrupt directional changes, as illustrated in Figure 5. In the unconditional block stacking task where the robot destroys the tower while stacking the last block, the tower-breaking transitions are highlighted. On the other hand, for successful trajectories on the second and third attribution maps, the attribution maps do not emphasize picking or stacking behaviors. Similarly, in the conditional block stacking task where the robot fails to stack the block, they spotlight the transitions of stacking behaviors.
To study the benefit of regularization on harder tasks, characterized by a larger trajectory space and a smaller fraction of the space observed in training, we explore  $\lambda$  values [0.0, 0.5, 1.0, 3.0, 5.0] while increasing the planning budget as illustrated in Figure 3. As the planning budget increases,  $\lambda = 0$  generates adversarial plans, resulting in decreased performance. In contrast, RGG+ demonstrates
Figure 5: Attribution maps for trajectories generated by diffusion planner highlight transitions that have a substantial contribution to the estimation of a high restoration gap by the gap predictor, indicated in red.
effectiveness across a wide range of  $\lambda$  values, with  $\lambda > 0$  consistently outperforming  $\lambda = 0$  (i.e., better than no regularization). Further investigation into the attribution method, perturbation magnitude, and comparison with guidance approaches, including metrics such as the rarity score, the negative realism score, and the discriminator, as well as the visualization of low and high restoration gap plans, can be found in Appendix B.
Metrics for Evaluating Generative Model Inception score (IS) (Salimans et al., 2016) and Fréchet inception distance (FID) (Heusel et al., 2017) are commonly used as standard evaluation metrics for generative models, assessing the quality of generated samples by comparing the discrepancy between real and generated samples in the feature space. However, these metrics do not distinguish between fidelity and diversity aspects of generated samples. To address this issue, precision and recall variants (Sajjadi et al., 2018; Kynkänniemi et al., 2019) are introduced to separately evaluate these properties. Subsequently, density and coverage (Naeem et al., 2020) are proposed to overcome some of the drawbacks of precision and recall, such as vulnerability to outliers and computational inefficiency. While these metrics are helpful for evaluating the quality of a set of generated samples, they are not suitable for ranking individual samples. In contrast, realism score (Kynkänniemi et al., 2019) and rarity score (Han et al., 2023) offer a continuous extension of improved precision and recall, enabling the evaluation of individually generated sample quality. Despite their usefulness,
these methods come with limitations as they rely on real samples for precise real manifold estimation, whereas our restoration gap does not have such a constraint.
Diffusion Model in Reinforcement Learning Diffusion models have gained prominence as a notable class of generative models, characterizing the data generation process through iterative denoising procedure (Sohl-Dickstein et al., 2015; Ho et al., 2020). This denoising procedure can be viewed as a way to parameterize the gradients of the data distribution (Song & Ermon, 2019), linking diffusion models to score matching (Hyvärinen & Dayan, 2005) and energy-based models (EBMs) (LeCun et al., 2006; Du & Mordatch, 2019; Nijkamp et al., 2019; Grathwohl et al., 2020). Recently, diffusion models have been successfully applied to various control tasks (Janner et al., 2022; Urain et al., 2023; Ajay et al., 2023; Chi et al., 2023; Liang et al., 2023). In particular, Diffuser (Janner et al., 2022) employs an unconditional diffusion model to generate trajectories consisting of state-action pairs. The approach includes training a separate model that predicts the cumulative rewards of noisy trajectory samples, which then guides the reverse diffusion process towards high-return trajectory samples in the inference phase, analogous to classifier-guided sampling (Dhariwal & Nichol, 2021). Building upon this, Decision Diffuser (Ajay et al., 2023) extends the capabilities of Diffuser by adopting a conditional diffusion model with reward or constraint guidance to effectively satisfy constraints, compose skills, and maximize return. Meanwhile, AdaptDiffuser (Liang et al., 2023) enhances generalization ability of the diffusion model to unseen tasks by selectively fine-tuning it with high-quality data, derived through the use of hand-designed reward functions and an inverse dynamics model. In contrast, in this work, we focus on evaluating the quality of individually generated samples and explore ways to enhance planning performance by utilizing guidance derived from these evaluations.
Restoring Artifacts in Generative Models Recently, several studies have concentrated on investigating the artifacts in Generative Adversarial Networks (GAN) model architectures for image generation tasks. GAN Dissection (Bau et al., 2019) explores the internal mechanisms of GANs, focusing on the identification and removal of units that contribute to artifact production, leading to more realistic outputs. In a subsequent study, an external classifier is trained to identify regions of low visual fidelity in individual generations and to detect internal units associated with those regions (Tousi et al., 2021). Alternatively, artifact correction through latent code manipulation based on a binary linear classifier is proposed (Shen et al., 2020). Although these methods can assess the fidelity of individual samples, they still necessitate additional supervision, such as human annotation. To address this limitation, subsequent works explore unsupervised approaches for detecting and correcting artifact generations by examining local activation (Jeong et al., 2022) and activation frequency (Choi et al., 2022). In contrast, our work primarily focuses on refining the generative process of diffusion probabilistic models to restore low-quality plans.
We have presented a novel refining method that fixes infeasible transitions within the trajectory generated by the diffusion planner. This refining process is guided by a proposed metric, restoration gap, which quantifies the restorability of a given plan. Under specific regularity conditions, we prove that the restoration gap effectively identifies unreliable plans while ensuring a low error probability for both type I and type II errors. The experimental results, which include enhancement in quantitative planning performance and visualization of qualitative attribution maps, highlight the importance of the refinement method of the diffusion planner.
Limitations While the restoration gap guidance effectively enhances the feasibility of plans and consistently improves the planning performance of diffusion models, our method is limited in situations where an offline dataset is provided. Training the diffusion model often requires transition data that uniformly covers the state-action space, the collection of which is a nontrivial and time-consuming task.
Future Work Our analysis of the effectiveness of the restoration gap is currently confined to a relatively simple task, Maze2D (see Figure 2), where we explicitly define normal and artifact plans. The choice of Maze2D is motivated by its suitability for identifying violations of prior knowledge, such as feasible plans not passing through walls. However, as future work, it would be worthwhile to explore the efficacy of restoration gap in more complex tasks, such as the block stacking task.
This work was supported by the Industry Core Technology Development Project, 20005062, Development of Artificial Intelligence Robot Autonomous Navigation Technology for Agile Movement in Crowded Space, funded by the Ministry of Trade, Industry & Energy (MOTIE, Republic of Korea) and by Institute of Information & communications Technology Planning & Evaluation (IITP) grant funded by the Korea government (MSIT) (No. 2022-0-00984, Development of Artificial Intelligence Technology for Personalized Plug-and-Play Explanation and Verification of Explanation, No.2019-0-00075, Artificial Intelligence Graduate School Program (KAIST)).
Ajay, A., Du, Y., Gupta, A., Tenenbaum, J., Jaakkola, T., and Agrawal, P. Is conditional generative modeling all you need for decision making? In International Conference on Learning Representations (ICLR), 2023.
Argenson, A. and Dulac-Arnold, G. Model-based offline planning. In International Conference on Learning Representations (ICLR), 2021.
Bau, D., Zhu, J.-Y., Strobelt, H., Zhou, B., Tenenbaum, J. B., Freeman, W. T., and Torralba, A. Gan dissection: Visualizing and understanding generative adversarial networks. In International Conference on Learning Representations (ICLR), 2019.
Beechey, D., Smith, T. M., and Şimşek, Ö. Explaining reinforcement learning with shapley values. In International Conference on Machine Learning (ICML), 2023.
Bertoin, D., Zouitine, A., Zouitine, M., and Rachelson, E. Look where you look! saliency-guided q-networks for generalization in visual reinforcement learning. In Advances in Neural Information Processing Systems (NeurIPS), 2022.
Chen, L., Lu, K., Rajeswaran, A., Lee, K., Grover, A., Laskin, M., Abbeel, P., Srinivas, A., and Mordatch, I. Decision transformer: Reinforcement learning via sequence modeling. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Chi, C., Feng, S., Du, Y., Xu, Z., Cousineau, E., Burchfiel, B., and Song, S. Diffusion policy: Visuomotor policy learning via action diffusion. In Proceedings of Robotics: Science and Systems (RSS), 2023.
Choi, H., Chang, W., and Choi, J. Can we find neurons that cause unrealistic images in deep generative networks? In International Joint Conference on Artificial Intelligence (IJCAI), 2022.
Deisenroth, M. and Rasmussen, C. E. Pilco: A model-based and data-efficient approach to policy search. In International Conference on Machine Learning (ICML), 2011.
Dhariwal, P. and Nichol, A. Diffusion models beat gans on image synthesis. In Advances in Neural Information Processing Systems (NeurIPS), 2021.
Du, Y. and Mordatch, I. Implicit generation and modeling with energy based models. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Fu, J., Kumar, A., Nachum, O., Tucker, G., and Levine, S. D4rl: Datasets for deep data-driven reinforcement learning. arXiv preprint arXiv:2004.07219, 2020.
Fujimoto, S., Meger, D., and Precup, D. Off-policy deep reinforcement learning without exploration. In International Conference on Machine Learning (ICML), 2019.
Garrett, C. R., Lozano-Pérez, T., and Kaelbling, L. P. Pddlstream: Integrating symbolic planners and blackbox samplers via optimistic adaptive planning. In Proceedings of the International Conference on Automated Planning and Scheduling (ICAPS), volume 30, pp. 440-448, 2020.
Grathwohl, W., Wang, K.-C., Jacobsen, J.-H., Duvenaud, D., and Zemel, R. Learning the stein discrepancy for training and evaluating energy-based models without sampling. In International Conference on Machine Learning (ICML), 2020.
Shen, Y., Gu, J., Tang, X., and Zhou, B. Interpreting the latent space of gans for semantic face editing. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2020.
Shrikumar, A., Greenside, P., and Kundaje, A. Learning important features through propagating activation differences. In International Conference on Machine Learning (ICML), 2017.
Silver, D., Huang, A., Maddison, C. J., Guez, A., Sifre, L., Van Den Driessche, G., Schrittwieser, J., Antonoglou, I., Panneershelvam, V., Lanctot, M., et al. Mastering the game of go with deep neural networks and tree search. nature, 529(7587):484-489, 2016.
Silver, D., Schrittwieser, J., Simonyan, K., Antonoglou, I., Huang, A., Guez, A., Hubert, T., Baker, L., Lai, M., Bolton, A., et al. Mastering the game of go without human knowledge. nature, 550 (7676):354-359, 2017.
Simonyan, K., Vedaldi, A., and Zisserman, A. Deep inside convolutional networks: Visualising image classification models and saliency maps. arXiv preprint arXiv:1312.6034, 2013.
Sohl-Dickstein, J., Weiss, E., Maheswaranathan, N., and Ganguli, S. Deep unsupervised learning using nonequilibrium thermodynamics. In International Conference on Machine Learning (ICML), 2015.
Song, Y. and Ermon, S. Generative modeling by estimating gradients of the data distribution. In Advances in Neural Information Processing Systems (NeurIPS), 2019.
Song, Y., Sohl-Dickstein, J., Kingma, D. P., Kumar, A., Ermon, S., and Poole, B. Score-based generative modeling through stochastic differential equations. In International Conference on Learning Representations (ICLR), 2021.
Talvitie, E. Model regularization for stable sample rollouts. In Proceedings of the Conference on Uncertainty in Artificial Intelligence (UAI), 2014.
Tassa, Y., Erez, T., and Todorov, E. Synthesis and stabilization of complex behaviors through online trajectory optimization. In 2012 IEEE/RSJ International Conference on Intelligent Robots and Systems, pp. 4906-4913, 2012.
Tousi, A., Jeong, H., Han, J., Choi, H., and Choi, J. Automatic correction of internal units in generative neural networks. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition (CVPR), 2021.
Urain, J., Funk, N., Peters, J., and Chalvatzaki, G. Se(3)-diffusionfields: Learning smooth cost functions for joint grasp and motion optimization through diffusion. In 2023 IEEE International Conference on Robotics and Automation (ICRA), pp. 5923-5930, 2023.
Voelcker, C., Liao, V., Garg, A., and Farahmand, A.-m. Value gradient weighted model-based reinforcement learning. In International Conference on Learning Representations (ICLR), 2022.
Williams, G., Aldrich, A., and Theodorou, E. Model predictive path integral control using covariance variable importance sampling. arXiv preprint arXiv:1509.01149, 2015.
Yu, T., Thomas, G., Yu, L., Ermon, S., Zou, J. Y., Levine, S., Finn, C., and Ma, T. Mopo: Model-based offline policy optimization. In Advances in Neural Information Processing Systems (NeurIPS), 2020.
Appendix A Proofs