SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Figure 3: Quantitative analysis of the sparsified models for open-loop prediction and planning. (a) Long-term future prediction error, with the shaded area representing the region between the  $25^{\text{th}}$  and  $75^{\text{th}}$  percentiles. The significant overlap between the curves suggests that reducing the number of ReLUs only leads to a minimal decrease in prediction accuracy. (b) Results of the trajectory optimization problem from Equation 10. We compare two optimization formulations: mixed-integer programming (MIP) and model-predictive path integral (MPPI), using models with varying levels of sparsification. The figure clearly indicates that MIP consistently outperforms its sampling-based counterpart, MPPI.
These evaluation environments make use of different modeling classes, including vanilla MLPs and complex GNNs. For closed-loop control evaluation, we additionally present the performance of our framework on two standardized benchmark environments from OpenAI Gym [7], CartPole-v1 and Reacher-v4.
- Piecewise affine function. We consider manually designed two-dimensional piecewise affine functions consisting of four linear pieces (Figure 2), and the goal is to recover the ground-truth system from data through the sparsification process starting from an overparameterized MLP. To train the model, we collect 1,600 transition pairs from the ground truth functions uniformly distributed over the 2D input space.
- Object pushing. A fully-actuated pusher interacts with an object moving on a 2D plane, as depicted in Figure 4a. The goal is to manipulate the object to reach a randomly generated target pose. We generated 50,000 transition pairs using the Pymunk simulator [6]. The observation  $y_{t}$  is defined by the position of four selected keypoints on the object, and the dynamics model is also instantiated as an MLP.
- Object sorting. In Figure 4c, a pusher is used to sort a cluster of objects that lie on a table into corresponding target regions. In this environment, we generate a dataset consisting of 150,000 transition pairs with two objects using Pymunk. Following the success of previous graph-based dynamics models [4, 37, 36], we use GNNs as the model class. The model takes the object positions as input and allows compositional generalization to extrapolate settings containing more objects, supporting up to 8 objects as tested in our benchmark.
- Rope manipulation. Figure 4b shows the task of manipulating a deformable rope into a target shape. We generate a dataset of 60,000 transition pairs through random interactions using Nvidia FleX [45]. We use an MLP to model the dynamics, and the observation  $y_{t}$  is the position of four selected keypoints on the rope.
4.1 How does the varying number of ReLUs affect the prediction accuracy?
Recover the ground truth piecewise affine system from data. The sparsification procedure starts with the full model with four hidden layers and 576 ReLU units. It then undergoes seven iterations of sparsification, with the number of remaining ReLUs, represented as  $\varepsilon_{k}$ , diminishing from 25 down to 2. As illustrated in Figure 2, the sparsified model, which retains only two ReLUs, accurately identifies the region partition and achieves a nearly zero distance from the ground truth. This enables
Figure 4: Qualitative results on closed-loop feedback control. (a) In object pushing, the objective is to manipulate the object to reach a randomly generated target pose, depicted as transparent in the first column. The second column illustrates how the planner, using the sparsified model, can leverage feedback from the environment to compensate for the modeling errors and accurately achieve the target. (b) The framework is also applicable to rope manipulation. Our sparsified model, in conjunction with the MIP formulation, facilitates closed-loop feedback control to manipulate the rope into desired configurations. (c) Our framework also consistently succeeds in object sorting tasks that involve complex contact events. Using the same model with the MIP formulation, the system can manipulate up to eight objects, sorting them into their respective regions.
the model to recover the underlying ground truth system and demonstrates the effectiveness of the sparsification procedure.
Future prediction using sparsified models at different sparsification levels. Existing literature provides comprehensive studies indicating that neural networks are overparameterized [22, 23, 16]. Still, we are interested in understanding how the proposed sparsification process affects the model prediction accuracy. We evaluate our framework on three benchmark environments, object pushing, sorting, and rope manipulation. Starting with the full model, we gradually sparsify it using decreasing values of  $\varepsilon_{k}$ . During training, we focus solely on the accuracy of one-step prediction but evaluate the models for their long-horizon predictive capability.
Figure 3a illustrates the prediction accuracy for models with varying numbers of ReLUs. "Object Sorting-2" denotes the task of sorting objects into two piles, while "Object Sorting-3" represents sorting into three piles. The blue curve shows the full-model performance, and the shaded area denotes the region between the  $25^{\text{th}}$  and  $75^{\text{th}}$  percentiles over 100 trials. The figure suggests that, even with significantly fewer ReLUs, the model still achieves a reasonable future prediction performance, with the confidence region significantly overlapping with that of the full model. It is worth
Figure 5: Quantitative analysis of model sparsification vs. closed-loop control performance. The horizontal axis represents the number of ReLUs remaining in the model, and the vertical axis indicates the closed-loop control performance. As shown in the figure, there exists a nice trade-off between the levels of model sparsification and the performance of closed-loop control. Models with fewer ReLUs are typically less accurate than the full model but make the MIP formulation tractable to solve. Across the spectrum of models, there exists a sweet spot, where a model, although only reasonably accurate, benefits from more powerful optimization tools and can lead to superior closed-loop control results. Moreover, our method consistently outperforms commonly used RL techniques such as PPO and SAC.
noting that our framework is adaptable to both vanilla MLPs (utilized in object pushing and rope manipulation) and GNNs (employed for object sorting), thereby showcasing the broad applicability of our proposed method. Later in Section 4.2 and 4.3, we will demonstrate that the sparsified models, although slightly less accurate than the full model, can yield superior open-loop and closed-loop optimization results when paired with more effective optimization tools.
4.2 How does the varying number of ReLUs affect open-loop planning?
Having obtained the sparsified models and examined their prediction accuracy, we next assess how these models can be applied to solve the trajectory optimization problem in Equation 10. The sparsified model contains significantly fewer ReLUs, making it feasible to use formulations with better optimality guarantees, as discussed in Section 3.5. Specifically, we formulate the optimization problem using MIP (Equation 11) and solve the problem using a commercial optimization solver, Gurobi [18]. We compare our method with MPPI, a commonly-used sampling-based alternative from the model-based RL community. As illustrated in Figure 3b, the MIP formulation permits the use of advanced branch-and-bound optimization procedures. With a sufficiently small number of ReLU units remaining in the neural dynamics models, we can solve the problem optimally. This consistently outperforms MPPI by a significant margin.
4.3 Can the sparsified model deliver better closed-loop control results?
The results in Section 4.2 only tell us how good different optimization procedures are as measured by the learned dynamics model. However, what we really care about is the performance when executing optimized plans in the original simulator or the real world. Therefore, it is crucial to evaluate the effectiveness of these models within a closed-loop control framework. Here we employ an MPC
$^{\ddagger}$ We omit the result of PPO on rope manipulation due to compute limitations, because our rope simulator does not support accelerated-time simulation and takes excessively long before PPO gains reasonable performance. We omit the result of SAC on CartPole-v1 because the Stable Baselines 3 SAC implementation does not support a discrete action space.
framework that, taking into account the feedback from the environment, allows the agent to make online adjustments to the action sequence.
Figure 4 visualizes multiple execution trials of object pushing, sorting, and rope manipulation in the real world using our method. Our framework reliably pushes the object to its target pose, deforms the rope into the desired shape, and sorts the many objects into the corresponding piles. We then present the quantitative results for object pushing, sorting, and rope manipulation, along with two tasks from OpenAI Gym [7], CartPole-v1 and Reacher-v4, measured in simulation, in Figure 5. Across various tasks, we observe a similar trade-off between the levels of model sparsification and closed-loop control performance. As the number of RLUs decreases, there is typically a slight decrease in prediction accuracy, but as illustrated in Figure 5, this allows us to formulate the trajectory optimization problem as an MIP and solve it using efficient branch-and-bound algorithms. Consequently, within the spectrum of sparsified models, there exists an optimal point where a model, albeit only reasonably accurate, benefits from the more effective optimization tools and can result in better closed-loop control performance. Our iterative sparsification process, discussed in Section 3.4, enables us to easily identify such model. Furthermore, our method consistently outperforms commonly used RL techniques such as PPO [59] and SAC [19] when using the same number of interactions with the underlying environments.
Conclusion. In this work, we propose to sparsify neural dynamics models for more effective closed-loop, model-based planning and control. Our formulation allows an end-to-end optimization of both the model class and the weight parameters. The sparsified models enable the use of efficient branch-and-bound algorithms and can deliver better performance in closed-loop control. Our framework applies to various dynamical systems and multiple neural network architectures, including vanilla MLPs and complicated GNNs. We also demonstrate the effectiveness and applicability of our method through its application to simple piecewise affine systems and manipulation tasks involving complex contact dynamics and deformable objects.
Our work draws inspiration and merges techniques from both the learning and control communities, which we hope can spur future investigations in this interdisciplinary direction to take advantage and make novel use of the powerful tools from both communities.
Limitations and future work. Our method relies on sparsifying neural dynamics models to fewer ReLU units to make the control optimization process solvable in a reasonable time due to the worst-case exponential run time of MIP solvers. Although our experiments showed that this already enabled us to complete a wide variety of tasks, our approach may struggle when facing a much larger neural dynamics model.
Our experiments also demonstrated superior closed-loop control performance using sparsified dynamics models with only reasonably good prediction accuracy as a result of benefiting from stronger optimization tools, but our approach may suffer if the sparsified dynamics model becomes significantly worse and incapable of providing useful forward predictions.
Acknowledgments. This work is in part supported by ONR MURI N00014-22-1-2740. Ziang Liu is supported by the Siebel Scholars program.
[1] Pulkit Agrawal, Ashvin Nair, Pieter Abbeel, Jitendra Malik, and Sergey Levine. Learning to poke by poking: Experiential learning of intuitive physics. arXiv preprint arXiv:1606.07419, 2016.
[2] Ross Anderson, Joey Huchette, Will Ma, Christian Tjandraatmadja, and Juan Pablo Vielma. Strong mixed-integer programming formulations for trained neural networks. Mathematical Programming, pages 1–37, 2020.
[3] Sajid Anwar, Kyuyeon Hwang, and Wonyong Sung. Structured pruning of deep convolutional neural networks. ACM Journal on Emerging Technologies in Computing Systems (JETC), 13 (3):1-18, 2017.
[60] Haochen Shi, Huazhe Xu, Zhiao Huang, Yunzhu Li, and Jiajun Wu. Robocraft: Learning to see, simulate, and shape elasto-plastic objects with graph networks. arXiv preprint arXiv:2205.02909, 2022.
[61] Kin Cheong Sou, Alexandre Megretski, and Luca Daniel. A quasi-convex optimization approach to parameterized model order reduction. In Proceedings of the 42nd annual Design Automation Conference, pages 933–938, 2005.
[62] HJ Suh and Russ Tedrake. The surprising effectiveness of linear models for visual foresight in object pile manipulation. arXiv preprint arXiv:2002.09093, 2020.
[63] Vincent Tjeng, Kai Xiao, and Russ Tedrake. Evaluating robustness of neural networks with mixed integer programming. arXiv preprint arXiv:1711.07356, 2017.
[64] Tianzhe Wang, Kuan Wang, Han Cai, Ji Lin, Zhijian Liu, Hanrui Wang, Yujun Lin, and Song Han. Apq: Joint search for network architecture, pruning and quantization policy. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 2078-2087, 2020.
[65] Manuel Watter, Jost Tobias Springenberg, Joschka Boedecker, and Martin Riedmiller. Embed to control: A locally linear latent dynamics model for control from raw images. arXiv preprint arXiv:1506.07365, 2015.
[66] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity in deep neural networks. arXiv preprint arXiv:1608.03665, 2016.
[67] Grady Williams, Andrew Aldrich, and Evangelos A Theodorou. Model predictive path integral control: From theory to parallel computation. Journal of Guidance, Control, and Dynamics, 40(2):344-357, 2017.
[68] Bichen Wu, Xiaoliang Dai, Peizhao Zhang, Yanghan Wang, Fei Sun, Yiming Wu, Yuandong Tian, Peter Vajda, Yangqing Jia, and Kurt Keutzer. Fbnet: Hardware-aware efficient convnet design via differentiable neural architecture search. In Proceedings of the IEEE/CVF Conference on Computer Vision and Pattern Recognition, pages 10734-10742, 2019.
[69] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Daydreamer: World models for physical robot learning. In Conference on Robot Learning, pages 2226-2240. PMLR, 2023.
[70] Kai Y Xiao, Vincent Tjeng, Nur Muhammad Shafiullah, and Aleksander Madry. Training for faster adversarial robustness verification via inducing relu stability. arXiv preprint arXiv:1809.03008, 2018.
[71] Lin Yen-Chen, Maria Bauza, and Phillip Isola. Experience-embedded visual foresight. In Conference on Robot Learning, pages 1015–1024. PMLR, 2020.
[72] Aojun Zhou, Yukun Ma, Junnan Zhu, Jianbo Liu, Zhijie Zhang, Kun Yuan, Wenxiu Sun, and Hongsheng Li. Learning n: M fine-grained structured sparse neural networks from scratch. arXiv preprint arXiv:2102.04010, 2021.
[73] Jiaji Zhou, Yifan Hou, and Matthew T Mason. Pushing revisited: Differential flatness, trajectory planning, and stabilization. The International Journal of Robotics Research, 38(12-13): 1477-1489, 2019.
[74] Barret Zoph and Quoc V Le. Neural architecture search with reinforcement learning. arXiv preprint arXiv:1611.01578, 2016.
A How does our method compare to prior works in model-based RL?
In this experiment, we aim to examine how the closed-loop control performance of our method compare to prior works in model-based reinforcement learning, evaluated on standard benchmark environments. We conduct experiments on two additional tasks from OpenAI Gym [7], CartPole-v1 and Reacher-v4, following the same procedures as described in Section 4.3. On top of a sampling-based planner (MPPI) and a model-free RL method (PPO), we employ two additional model-based RL methods, (1) using PPO to learn a control policy from our learned full neural dynamics model, and (2) MBPO [29] learning a model and a policy from scratch. The model-based RL methods require additional time to train a policy using the learned dynamics model, whereas our approach directly optimizes a task objective over the dynamics model without needing additional training.
The experiment results shown in Figure 6 further demonstrate the superior performance of our approach compared to prior methods on the two standard benchmark tasks. Notably, our approach achieves better performance with highly sparsified neural dynamics models with fewer RLUs compared to prior works.
Figure 6: Closed-loop control performance of our method (MIP) compared against prior methods on two new environments. Our method with fewer ReLUs outperforms prior methods using models with more ReLUs, and we similarly observe a sweet spot that balances between model prediction accuracy and control performance.
Figure 7: We tested the closed-loop control performance of dynamics models trained and simplified using our method by incorporating them as the forward model in a model-based RL framework optimized with PPO. Our findings indicate that even when the dynamics models are substantially simplified, they continue to allow for satisfactory control performance.
B Do models trained using our approach generalize to prior model-based RL methods?
The neural dynamics model learned in our method is generic and not limited to only working with our planning framework. We take the learned full and sparsified dynamics models on the CartPole-v1 environment and train a control policy with PPO interacting only with the learned model, and provide the experiment results in Figure 7.
(a) Piecewise affine function #1
Figure 8: Comparison with the channel pruning baseline on recovering the PWA functions. We evaluate both our sparsification pipeline and the channel pruning baseline using two hand-designed piecewise affine functions, each composed of four linear pieces. Our pipeline successfully recovers the underlying ground truth system, whereas the baseline does not provide an accurate match.
(b) Piecewise affine function #2
The results show that the neural dynamics models trained in our method can generalize and combine with another model-based control framework. As the model becomes progressively sparsified, the closed-loop control performance gracefully degrades.
C How does our sparsification technique compare to prior neural network pruning methods?
The aim of this study is to discuss and compare our sparsification technique with the pruning methods commonly employed in the field. Most pruning strategies in existing literature primarily focus on eliminating as many neurons as possible to reduce computation and enhance efficiency. In contrast, our method aims to eliminate as many nonlinearities from the neural networks as possible. This differs from channel pruning, which only zeroes out values. Our approach permits the replacement of ReLU activations with identity mappings, the inclusion of which allows a more accurate model to be achieved at an equivalent level of sparsification. This offers a considerable advantage during the planning stage.
To illustrate our point more concretely, we provide, in this section, experimental results comparing our method against Level 1 Channel Pruning as referenced in [35].
C.1 Evaluation on Piecewise Affine (PWA) Functions
The two pruning methods are tasked to recover the ground truth PWA functions, as detailed in the experiment section of the main paper. Figure 8 illustrates the results after sparsifying the neural networks to two rectified linear units (ReLU) using both methods. Our method successfully identifies the region partition and the correct equations for describing values of each region, whereas the baseline [35] exhibits noticeable deviations.
C.2 Evaluation on Dynamics Prediction
In this section, we extend the comparison to three other tasks: object pushing, object sorting, and rope manipulation. For the object pushing and rope manipulation tasks, we train the neural dynamics model for a defined number of epochs before pruning is carried out by masking particular channels. Post-pruning, model speedup is performed using Neural Network Intelligence Library [51] to alter the model's architecture and remove masked neurons. This process is repeated as further channels are pruned and the models are finetuned for additional epochs.
For the object sorting task involving graph neural networks, we perform a similar procedure to construct the baseline. During the initial model training phase, the mask resulting from the  $L_{1}$  norm pruning operation is used to nullify specific weights, and the corresponding gradients are also masked during the finetuning phase.
To ensure fairness and reliability in the comparison, we maintain identical settings for both our sparsification technique and the pruning baseline. Therefore, for every round of compression, both models are subjected to the same number of training epochs, using the same dataset, and are reduced to the same number of ReLU units.
Figure 9: Dynamics prediction error of sparsified models using our method vs. baseline. We compare the dynamics prediction error of models sparsified using our method against models sparsified using the channel pruning method proposed by Li et al. [35]. The x-axis represents the number of remaining ReLU units in the model. The y-axis represents the prediction error measured by the root mean squared error between the prediction and the ground truth next state. Because our sparsification method only targets non-linear units while allowing linear units, models sparsified using our method constantly exhibit lower prediction error across all task settings.
We provide quantitative comparisons between our sparsification method and the baseline in Figure 9. Throughout the sparsification process, because our sparsification objective allows replacing non-linearities with identity mappings, our method consistently achieves a superior performance measured by prediction error, across all tasks.
C.3 Evaluation on Closed-Loop Control
In this experiment, we aim to further examine whether our method also boosts the closed-loop control performance that is critical for executing the optimized plans in the real world. We choose the object pushing task, and prune the learned dynamics model down to 36, 24, 18, and 15 ReLU units using our proposed sparsification method and the channel pruning method proposed by Li et al. [35] respectively. As shown in Figure 10, models pruned using our method consistently exhibit superior closed-loop control performance across all sparsification levels.
Figure 10: An ablation of our sparsification method compared with a prior network pruning method, evaluated by closed-loop control performance, demonstrating superior performance in closed-loop control.