Model-Based Control with Sparse Neural Dynamics
Authors: Ziang Liu
Genggeng Zhou
Jeff He
Tobia Marcucci
Li Fei-Fei
Jiajun Wu
Yunzhu Li
Venue: NeurIPS 2023 poster (2023)
Keywords: model learning
model-based control
neural network sparsification
mixed-integer programming
trajectory optimization
TL;DR: We propose to sparsify neural dynamics models for more effective closed-loop, model-based planning using the mixed-integer formulation.

Abstract: Learning predictive models from observations using deep neural networks (DNNs) is a promising new approach to many real-world planning and control problems. However, common DNNs are too unstructured for effective planning, and current control methods typically rely on extensive sampling or local gradient descent. In this paper, we propose a new framework for integrated model learning and predictive control that is amenable to efficient optimization algorithms. Specifically, we start with a ReLU neural model of the system dynamics and, with minimal losses in prediction accuracy, we gradually sparsify it by removing redundant neurons. This discrete sparsification process is approximated as a continuous problem, enabling an end-to-end optimization of both the model architecture and the weight parameters. The sparsified model is subsequently used by a mixed-integer predictive controller, which represents the neuron activations as binary variables and employs efficient branch-and-bound algorithms. Our framework is applicable to a wide variety of DNNs, from simple multilayer perceptrons to complex graph neural dynamics. It can efficiently handle tasks involving complicated contact dynamics, such as object pushing, compositional object sorting, and manipulation of deformable objects. Numerical and hardware experiments show that, despite the aggressive sparsification, our framework can deliver better closed-loop performance than existing state-of-the-art methods.

Model-Based Control with Sparse Neural Dynamics
Ziang Liu .2
Genggeng Zhou²*
Jeff He
Tobia Marcucci
Li Fei-Fei
Jiajun Wu? Yunzhu Li²4
CCornlll University Stanford University
3 Massachusetts Institute of Technology
UUniversity of Illinois Urbana-Champann
ziangliu@cs cornell. edu fg9zhou, jwf.eee4) @stanfrdd edu tobiam@mit.co edu
beaifeli.jjajjunuuulll stanford. edu.
yunzhuli@illinois edu
Abstract
Learning predictive models from observations using deep neural net- works (DNNs) is a promising new approach to many real-world planning and control problems. However, common DNNs are too unstructured for effective planning, and current control methods typically rely on extensive sampling or local gradient descent. In this paper. we propose new framework for integrated model learning and predictive control that is amenable to efficient optimization al- gorithms. Specifically, start with ReLU neural model of the system dynam- ics and, with minimal losses in prediction accuracy, we gradually sparsify it by removing redundant neurons. This discrete sparsification process is approximated as a continuous problem, enabling an end-to-end optimization of both the model architecture and the weight parameters. The sparsified model subsequently used by a mixed-integer predictive controller, which represents the neuron acti- vations binary variables and employs efficient branch-and-boune algorithms. Our framework is applicable to wide variety of DNNs, from simple multilayer perceptrons complex graph neural dynamics. can efficiently handle tasks involving complicated contact dynamics, such as object pushing, compositional object sorting, and manipulation of deformable objects. Numerical and hardware experiments show that, despite the aggressive sparsification, our framework can deliver better closed-loop performance than existing state-of-the-ar methods. f
Introduction
Our mental model of the physical environment enables us to easily carry out a broad spectrum of complex control tasks, many of which lie far beyond the capabilities of present-day robots [32]. is, therefore, desirable to build predictive models of the environment from observations and de- velop optimization algorithms to help the robots understand the impact of their actions and make effective plans to achieve a given goal. Physics-based models [26, 73] have excellent generalization ability but typically require full-state information of the environment, which hard and sometimes impossible to obtain in complicated robotic (manipulation) tasks. Learning- eling cireumvents the problem by learning a predictive model directly from raw sensory observa- tions, and recent successes are rooted in the use of deep neural networks (DNNs) as the functional class [14, 21, 56, 47]. Despite their improved prediction accuracy, DNNs are highly nonlinear,
dynamics mod-
denotes equal contribution
Please see our website
for additional visualizations.
37th Conference on Neural Information Processing Systems (NcurIPS 2023).

---

Observation
Action
Observation
Action
ReLL
ReLLU
ReLU
RoLU
Sparsify
RaLL
RaLLU
ReLU
RoLL
ID
RoLLU
Piecewise Affine Systems
Object Pushing
RoLL
RoLU
HaLL
HoLL
ID
New Observation
New Observation
(a) Sparsify the learned neural dynamics
Observation
Observation, Action
World
MIP Planner
Neural Dynamics
Action
New Observation
Object Sorting
Rope Manipulation
(b) Planning with sparsified neural dynamics
(c) Benchmark tasks
Figure 1: Model-based control with sparse neural dynamies. (a) Our framework sparsifies the neural dy-
namies models by either removing neurons or replacing ReLU activation functions with identity mappings (ID), (b) The sparsified models cnable the use of efficient MIP methods for planning. which can achieve better closed- loop performance than sampling-based alternatives commonly used in model-based RL. (c) We evaluate our framework on various dynamical systems that involve complex contact dynamics, including tasks like object pushing and sorting, and manipulating deformable rope.
making model-based planning with neural dynamics models very challenging. Existing methods often rely on extensive sampling or local gradient descent to compute control signals, and can be ineffective for complicated and long-horizon planning tasks.
Compared to DNNs, simpler models like linear models are amenable to optimization tools with bet- ter guarantees, but often struggle to accurately fit observation data. An important question arises: how precise do these models need to when employed within a feedback control loop? The cogni- tive science community offers substantial evidence suggesting that humans do not maintain highly accurate mental models; nevertheless, these less precise models can be effectively used with environ- mental feedback [30, 10]. This notion is also key in control control-oriented system identification [25, 44] and model order reduction [55, 61]. The framework from Li et al. [38] trades model expressive- ness and precision for more efficient and effective iriimization-based planning through the learning of compositional Koopman operators. However, their approach is limited by the linearity of the representation in the Koopman embedding space and struggles with more complex dynamics.
In this paper. we propose a framework for integrated model learning and control that trades off prediction accuracy for the use of principled optimization tools. Drawing inspiration from the neural network pruning and neural architecture search communities [22, 54, 16, 5, 41], we start from a neural network with ReLU activation functions and gradually reduce the nonlinearity of the model by removing ReLU units or replacing them with identity mappings (Figure la). This yields a highly sparsified neural dynamics model, that is amenable to model-based control using tifte--f-the-att solvers for mixed-integer programming (MIP) (Figure 1b).
We present examples where the proposed sparsification pipeline can determine region partition and uncover the underlying system for simple piccewise affine systems. Moreover, can maintain high prediction accuracy for more complex manipulation tasks, using a considerably smaller portion of the original nonlinearities. Importantly, our approach allows the joint optimization of the network architecture and weight parameters. This yields a spectrum of models with varying degrees of sparsification. Within this spectrum, we can identify the simplest model that is adequate to meet the requirements of the downstream closed- loop control task.
Our contributions can summarized as follows: (i) We propose novel formulation for identifying the dynamics model from observation data. For this step, we introduce continuous approximation of the sparsification problem, enabling end-to-end gradient-hass based optimization for both the model class and the model parameters (Figure la). (ii) By having significantly fewer ReL.U units than

---

the full model, the sparsified dynamics model allows us to solve the predictive-contor problems using efficient MIP solvers (Figure lb). This can lead better closed-loop performance compared to both model-free and model-based reinforcement learning (RI.) baselines. (iii) Our framework can be applied to many types of neural dynamics, from vanilla multilayer perceptrons (MLPs) to complex graph neural networks (GNNs). We show its effectiveness in variety of simulated and real-world manipulation tasks with complex contact dynamics, such as object pushing and sorting, and manipulation of deformable objects (Figure Ic).
2
Related Work
Model learning for planning and control. Model-based RL agents learn predictive models of their environment from observations, which are subsequently used to plan their actions [9, 53]. Re- cent successes in this domain often heavily rely on DNNs, exhibiting remarkable planning and control results in challenging simulated tasks [58], as well as complex real-world locomotion and manipulation tasks [34, 56]. Many of these studies draw inspiration from advancements in computer vision, learning dynamics models directly in pixel-space [15, 11, 12, 71, 62), keypoint representation [31, 47, 391, particle/mesh representation [36, 60, 27]. or low-dimensional latent space [65, 1, 21, 20, 58, 69]. While previous works typically assume that the model class is given and fixed during the optimization process, our work puts emphasis on finding the desired model class via an aggressive network sparsification, to support optimization tools with better guarantees. We are willing to sacrifice the prediction accuracy for better closed-loop performance using more principled optimization techniques.
Network sparsification. The concept of neural network sparsification is not new and traces back to the 1990s [33]. Since then, extensive research has been conducted, falling broadly into two categories: network pruning [23, 22, 66, 54, 35, 24, 3, 43, 16, 42, 5, 72] and neural architecture search [74, 8, 40, 13, 64]. Many of these studies have demonstrated that fitting an overparameterized model before pruning yields better results than directly fitting smaller model. Our formulation is closely related to DARTS [41] and FBNet [68], which both seek a continuous approximation of the discrete search process. However, unlike typical structured network compression methods, which try to remove as many units as possible, our goal here is to minimize the model nonlinearity. To this end, method also permits the substitution of ReLU activations with identity mappings. This leaves the number of units unchanged but makes the downstream optimization problem much simpler.
Mixed-integer modeling of neural networks. The input-output map of neural network with ReLU activations is piecewise affine function that can be modeled exactly through a set of mixed- integer linear inequalities. This allows us to use highly-effectivv MIP solvers for the solution of the model-based control problem. The same observation has been leveraged before for robustness analysis of DNNs in [63, 70], while the efficiency of these mixed-integer models has been thoroughly studied in [2].
3 Method
In this section, we describe our methods for learning a dynamics model using environmental ob- servations and for sparsifying DNNs through continuous approximation of the discrete pruning process. Then we discuss how the sparsified model can be used by an MIP solver for trajectory optimization and closed-loop control.
3.1 Learning a dynamics model over the observation space
Assume we have a dataset D = {v,"" u.rh) | - T.m ,,..,,M,,'' M collected via interac- tions with the environment, where y/2" and u." denote the observation and action obtained time in trajectory m. Our goal is to learn an autoregressive model fe. parameterized by 6, as a proxy of the real dynamics that takes small sequence of observations and actions from time t' to the current time t, and predicts the next observation at time 1:
0 = fo v.a))
(I)
3

---

We optimize the parameter 0 to minimize the simulation error that describes the long-term discrep- ancy between the prediction and the actual observation:
M
(
llv+1
fo )
(2)
m===
3.2 Neural network sparsification by removing or replacing ReLU activations
We instantiate the transition function fa ReLU neural network with N hidden layers. Let us denote the number of neurons in the 7t layer as N₁. When given an ipput = (yP:t. uP:r)) ) we denote the value of the fth neuron in layer i before the ReLU activation as Vij- Regular ReLU neural networks apply the rectifier function to every Tij and obtain the activation value using ztj = (zy( Rx))( & max(0, The nonlinearity introduced by the ReLU function allows the neural networks to fit the dataset but makes the downstream planning and control tasks more challenging. As suggested by many prior works in the field of neural network compression [22, 16), a lot of these ReLUs are redundant and can be removed with minimal effects on the prediction accuracy. In this work, we reduce the number of ReLU functions by replacing the function hıs with either an identity mapping (D:r))) A Fij or a zero function Zero(rij) P 0, where the latter is equivalent to removing the neuron (Figure la).
We divide the parameters in fa into two vectors, = (w,a), The vector collects the weight matrices and the bias terms. The vector a consists of a set of integer variables that parameterize the architecture of the neural network: far E (1,2,3)
1, N.. ,,, NN]. such that
(ReLU(r))) (Z1g) = ID(r,s)
if aaj =
if aij = 2 if aij = 3
(3)
Zero(zis))
The sparsification problem can then be formulated the following MIP:
min (=
C(0)
S.L.
ΣΣ
1(ay 1) Se.
2=1==
where 11 is the indicator function, and the value of decides the number of regular ReLU functions that are allowed to remain in the neural network.
3.3 Reparameterizing the categorical distribution using Gumbel-Softmax
Solving the optimization problem in Equation is hard, as the number of integer variables in a equals the number of ReLU neurons in the neural network, which is typically very large. Therefore, we relax the problem by introducing a random variable тay indicating the categorical distribution of
ai assigning to one of the three categories, where Tzj reformulate the problem as:
dk p((as = k) for k = 1, 3. We can then
N
N
min W
E[C(O)]
s.t.
Σ
EE
aij ~ Tij
(5)
1=1
where π 2 [mi 1i.,,, N.j
N.).
In Equation 5, the sampling procedure 0nf m.y is not differentiable. In order to make end-to-end gradient-based optimization possible, we employ the Gumbel-Softman [28, 46) technique to obtain continuous approximation of the discrete distribution.
Specifically, for a 3-class categorical distribution Ti.. where the class probabilities are denoted as Tiys T Ty. Gumbel-Max [17] allows draw 3-dimensional one-hot categorical samples ăy from the distribution via:
Eig = OneHot(arg max , + )),
(6)
&

---

where g are i..d. samples drawn from (0, 1), which is obtained by sampling un ~ Uniform(O, 1) and computing ge log(- log( ( )). We can then use the softmax function as continuous, differentiable approximation the arg max function:
exp ((log 1 + )/T)
Ex by exp (foo m. + )/-)
We denote this operation as 20 ~ Concrete ( 7) [46], where т is a temperature parameter con- trolling how close the softmax approximation is to the discrete distribution. As the temperature approaches zero, samples from the Gumbel-Softmax distribution become one-hot and identical to the original categorical distribution.
After obtaining Zij. we can calculate the activation value ztty as a weighted sum of different fune- tional choices:
2 b Z/y ReLU((r))) + + (D(s))) + zly Zero(ray);
and then use gradient descent to optimize both the weight parameters w and the architecture distri- bution parameters
During training, we can also constrain žy to be one-hot vectors by using arg max, but use the continuous approximation in the backward pass by approximating Vaãi so Vozij. This is denoted as "Straight-Through" Gumbel Estimator [28].
3.4 Optimization algorithm
Instead of limiting the number of regular ReLUs from the very beginning of the training process, start with randomly initialized neural network and use gradient descent optimize and by minimizing the following objective function until convergence:
+R
(9)
where the regularization term R(x) Σ E my aims explicitly reduce the use of the regular ReLU function. One could also consider adjusting it to R(x) e Σ N; (nly ADMT)) with small ÅιD discourage the use of identity mappings at the same time.
We then take an iterative approach by starting with relatively large é and gradually decrease its value for K iterations with E > E2 >> ... > EK =. Within each optimization iteration using EE, we first rank the neurons according to max( m3y) in descending order, and assign the activation function for the top-ranked neurons as ID πミ Z T. or Zero otherwise, while keeping the bottom ER neurons intact using Gumbel-Softme optimizing w and ı using gradient descent to minimize Equation 9. The sparsification process generates range of models various sparsification levels for subsequent investigations.
described in Section 3.3. Subsequently,
continue
3.5 Closed-loop feedback control using the sparsified models
After we have obtained the sparsified dynamics models, we fix the model architecture and formulate the model-based planning task as the following trajectory optimization problem:
min Edlee u+)
s.t.
Me+1 fo (yera (yrra)
(10)
where e the cost function. When the transition function fo highly nonlinear neural network, solving the optimization problem is not easy. Previous methods [71, 11, 56, 14, 47] typically regard the transition function as a black box and rely on sampling based algorithms like the cross-entropy method (CEM) [57] and model -rredccpee path integral (MPPI) [67] for online planning. Others have also tried applying gradient descent to derive the action signals [36, 37J. However, the number of required samples grows exponentially with the number of inputs and trajectory length. Gradient descent can also be stuck local optima, and is also hard to assess the optimality or robustness of the derived action sequence using these methods.
5

---

Ground Truth
Full (576 ReLUs)
Sparsified (2 ReLUs)
Ground Truth
Full (576 ReLUs)
Sparsified (2 ReLUs)
(a) Piecewise affine function #1
(b) Piecewise affine function #2
Figure 2: Recover ground truth piccewise affine functions from data. We evaluate our sparsification pipeline on two hand designdd piecewise affine functions composed of four linear pieces. Our pipeline suc- cessfully generates sparsified models with 2 ReLUs that accurately fit the data, determine the region partition, and recover the underlying ground truth system.
3.5.1 Mixed-integer formulation of ReLU neural dynamics
The sparsified neural dynamics models open up the possibility of dissecting the model and solving the problem using more principled optimization tools. Specifically, given that a ReLU neural net- work is piecewise affine function, we can formulate Equation 10 MIP. We assign to each ReLU a binary variable
1 (x: 2 0) to indicate whether the pre-activation value is larger smaller than zero. Given lower and upper bounds on the input S S (which we calculate by passing the offline dataset through the sparsified neural networks), the equality x+ ReLU(()) å max(0, x) can be modeled through the following set of mixed-integer linear constraints:
xt S ((1 a), x+ 2 x, x² S ua,
xt 220, 0, a (0,
(11)
If only a few ReLUS are left in the model, Equation 10 can efficiently solved to global optimality.
The formulation in Equation 11 is the simplest mixed-integg encoding of ReLU network, and variety of strategies are available in the literature to accelerate the solution of our MIPs. For large- scale models, it is possible to warm start the optimization process using sampling-based methods or gradient descent, and subsequently refine the solution using MIP solvers [49]. There also exist more advanced techniques to formulate the MIP [2, 48, 50], these lead to tighter convex relaxations of our problem and allow us to identify high-quality solutions of Equation 10 carlier in the rranch-and-- bound process. The ability to find globally-optimal solutions is attractive but requires the model to exhibit a reasonable global performance. The sparsification step helps us also in this direction, since typically expect smaller simulation error from sparsified (simpler) model than its unsparsified (very complicated) counterpart when moving away from the training distribution. In addition, we could also explicitly counteract this issue with the addition of trust-region constraints that prevent the optimizer from exploiting model inaccuracies in the areas of the input space that are not well- supported by the training data [52].
3.5.2 Tradeoff between model accuracy and closed-loop control performance
Models with fewer ReLUs are generally less accurate but permit the use of more advanced opti- mization tools, like efficient branch-and-bound algorithms implemented in state-of-the-art solvers. Within a model-predictive control (MPC) framework, the controller can leverage the environmental feedback counteract prediction errors via online modifications of the action sequence. The itera- tive optimization procedure in Section 3.4 yields a series of models at different sparsification levels. By comparing their performances and investigating the trade-off between prediction accuracy and closed-loop control performance, we can select the model with the most desirable capacity.
Experiments
In our experiments, we seek to address three central questions: (1) How does the varying number of ReLUs affect the prediction accuracy? (2) How does the varying number of ReLUs affect open- loop planning? (3) Can the sparsified model, when combined with more principled optimization methods, deliver better closed-loop control results?
Environments, tasks, and model classes. We evaluate our framework on four environments spec- ified in different observation spaces, including state, keypoints, and object-centric representations.

---

Object Pushing
Object Sorting-2
Object Sorting-3
Rope Manipulation
1.6
512
768
60
36
36
24
24
24
24
RMSE
15
15
16
Prediction
0.4
12
0.4
12
0.5-
12
04
2.
0.2-
0.3
0.0
To
0.00
0.0
10
0.0g
Time Stop
Time Stop
Time Step
(b) Trajectory Optimization
MIP
MPPI
Objective
1.5-
10
1.
5,5
24 Remaining ReLUs
15
36
12
24
Remaining ReLUs
Remaining ReLUss
Remaining RaLus
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 25th and 75 percentiles. The significant overlap between the curves suggests that reducing the number of ReLUS only leads to minimal decrease prediction accuracy. (b) Results
trajectory optimization problem from Equation 10. We compare two optimization formulations: mixed-intedd integer programming (MIP) and moodel-predictve 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 OpenAl Gym [7], CartPole-vl 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 randomly generated target pose. We generated 50,000 transition pairs using the Pymunk simulator [6]. The observation ye is defined by the position of four selected keypoints 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, generate 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 objects tested in our benchmark.
Rope manipulation. Figure 4b shows the task manipulating deformable rope into a target shape. We generate 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/: 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. then undergoes seven iterations of sparsification, with the number of remaining ReLUs, represented as Ex. diminishing from 25 down As illustrated Figure the sparsified model, which retains only two ReLUs, accurately identifies the region partition and achieves nearly zero distance from the ground truth. This enables
7

---

Time
Time
Trial
Intial
Target
Target
Trial #2
Intal
Initial
Targel
Target
(a) Object Pushing
(b) Rope Manipulation
Time
Trial #1
Trial #2
(c) Object Sorting
Figure 4: Qualitative results on closed-loop feedback control. (a) In object pushing, the objective is to manipulate the object to reach 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 also applicable to rope manipulation. Our sparsified model, in conjunction with the MIP formulation, facili- tates 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 manipulate up eight objects, sorting them into their respective regions.
the model 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, gradually sparsify it using decreas- ing values εg. 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 25th and 75th percentiles over 100 trials. The figure suggests that, even with significantly fewer ReL.Us, the model still achieves a reasonable future prediction perfor- mance, with the confidence region significantly overlapping with that of the full model. It is worth

---

Object Pushing
Reacher-v4
CartPole-v1
MIP
MPPI
0.25
MPPI
MIP
RMSE
0.6
SAC
PPO
RMSE
0.20
PPO
SAC
RMSE
0.009
MPPI
PPO
0.15
0.007
0.4
1
0.2
0.05
0.0
200
768
0.00
1624
0.003
0
16
48
Remaining ReLUs
Remaining ReLUs
Remaining ReLUs
Object Sorting-2 2
Object Sorting-3
Rope Manipulation
MIP
1.4
MIP
MIP
MPPI
MPP
0.60
MPPI
1.2
PPO
1.2
PPO
SAC
RMSE
1.0-
RMSE
SAC
RMSE
0.8-
0.8-
0.30
0.6
0.6
0.15
0.4
121524
60 100
121524
100
512
0.00
12:18
36 Remaining ReLUs
200
768
Remaining ReLUs
Remaining ReLUs
Figure 5: Quantitative analysis of model sparsification vs. closed-loop control performance. The horizon- tal axis represents the number of ReLUs remaining in the model, and the vertical axis indicates the closed-loop control performance.
shown the figure, there exists a nice trade-off between the levels of model sparsifi- cation and the performance closed-loop control. Models with fower ReLUs are typically less accurate than the full model but make the MIP formulation tractable to solve. Across the spectrum of models, there exists sweet spot, where model, although only reasonably accurate, benefits from more powerful optimization tools and can lead superior closed-loop control results. Moreover, our method consistently outperforms commonly used RI. techniques such as PPO and SAC.
noting that our framework is adaptable to both vanilla MLPs (utilized 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 mod- els, 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 spar- sified model contains significantly fewer ReLUs, making it feasible use formulations with bet- ter 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-base -ased alternative from model-based RL community. As illustrated in Figure 3b, the MIP formulation permits the use of advanced branch-and-boun
optimization procedures. With sufficiently small number of ReLU units remaining in the neural dynamics models, we can solve the problem optimally. This consistently outperforms MPPI by 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 these models within a closed-loop control framework. Here we employ an MPC
We omit the result 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 perfor- mance. We omit the result of SAC on CartPole-vi because the Stable Baselines SAC implementation does support 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 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 OpenAl Gym [7], CartPole-
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 ReLUs decreases, there is typically slight de- crease in prediction accuracy, but as illustrated in Figure 5, this allows to formulate the trajectory optimization problem as an MIP and solve it using efficient branch-an sanech-nnd-bunn algorithms. Conse- quently, within the spectrum of sparsified models, there exists an optimal point where model, albeit only reasonably accurate, benefits from the more effective optimization tools and can result in better closed-looy control performance. Our iterative sparsification process, discussed in Section 3.4, en- ables us 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.
5 Discussion
Conclusion. In this work, we propose to sparsify neural dynamics models for more effective closed- loop, model-based planning and control. Our formulation allows end-to-end optimization of both the model class and the weight parameters. The sparsified models enable the use of efficient branch-and boundooudddd algorithms and can deliver better performance in closed-loop control. Our frame- work applies to various dynamical systems and multiple neural network architectures, including vanilla MLPs and complicated GNNs. We also demonstrate the effectiveness and upplicability of our method through its application to simple piecewise affine systems and manipulation tasks in- volving 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 make the control optimization process solvable in reasonable time due the worst- case exponential run time of MIP solvers. Although our experiments showed that this already en- abled to complete wide variety of tasks, our approach may struggle when facing much larger neural dynamics model.
Our experiments also demonstrated superior closed-loop control performance using sparsified dy- namics models with only reasonably good prediction accuracy as a result optimization tools, but our approach may suffer if the sparsified dynamics model becomes signifi- cantly worse and incapable of providing useful forward predictions.
of benefiting from stronger
Acknowledgments.
This work is in part supported by ONR MURI N00014-22-1-2744 Ziang Liu supported by the Siebel Scholars program.
References
[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: 160666041999 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 Computing Systems (JETC), 13 (3): 1-), 2017.
10

---

[4] Peter Battaglia, Razvan Pascanu, Matthew Lai, Danilo Jimenez Rezende, et al. Interaction networks for learning about objects, relations and physics. Advances in neural information processing systems, 29, 2016.
[5] Davis Blalock, Jose Javier Gonzalez Ortiz, Jonathan Frankle, and John Guttag. What is the state of neural network pruning? arXiv preprint arXiv:2003.03033 2020.
[6] Victor Blomqvist. Pymunk. https://pymunk: org. November 2022. Version 6.4.0.
[7] Greg Brockman, Vicki Cheung, Ludwig Pettersson, Jonas Schneider, John Schulman, Jie Tang, and Wojciech Zaremba. Openai gym. CoRR, abs/1606.011400 2016. URL http://arxiv. arg/abs/ee heee/bbs/1666.... 1606.660600 01540.
[8] Han Cai, Ligeng Zhu, and Song Han. Proxylessnas: Direct neural architecture search on target task and hardware. arXiv preprint arXiv:7812.0033: 2018.
[9] Mare Deisenroth and Carl E Rasmussen. Pilco: A model-based and data-efficient approach policy search. Proceedings of the 28th International Conference on machine learning (ICML-11), pages 465 4722. Citeseer, 2011.
[10] James Doyle and David N Ford. Mental models concepts for system dynamics research. System dynamics review: the journal of the System Dynamics Society, 14(1):3-29, 1998.
[11] Frederik Ebert, Chelsea Finn, Alex X Lee, and Sergey Levine. Self-supervised visual planning with temporal skip connections. In CoRL, pages 344 444-356,20 2017.
[12] Frederik Ebert, Chelsea Finn, Sudeep Dasari, Annie Xie, Alex Lee, and Sergey Levine. Visual foresight: Model-based deep reinforcement learning for vision-based robotic control. arXiv preprint 1812.008100 2018.
[13] Thomas Elsken, Jan Hendrik Metzen, Frank Hutter, et al. Neural architecture search: survey.
J. Mach. Learn. Res., 20(55):1 2019.
[14] Chelsea Finn and Sergey Levine. Deep visual foresight for planning robot motion. In 2017 IEEE International Conference on Robotics and Automation (ICRA), pages 2786-2793, IEEE, 2017.
[15] Chelsea Finn, Ian Goodfellow, and Sergey Levine. Unsupervised learning for physical inter- action through video prediction. arXiv preprint arXiv:1605.0713: 071005 2016.
[16] Jonathan Frankle and Michael Carbin. The lottery ticket hypothesis: Finding sparse, trainable neural networks. arXiv preprint arXiv: 183.00660 2018.
[17] Emil Julius Gumbel. Statistical theory of extreme values and some practical applications: series of lectures, volume 33. US Government Printing Office, 1954.
[18] Gurobi Optimization, LLC. Gurobi Optimizer Reference Manual, 2023. URL https:///wwr gurobi com.
[19] Tuomas Haarnoja, Aurick Zhou, Pieter Abbeel, and Sergey Levine. Soft actor-critie: Off- policy maximum entropy deep reinforcement learning with a stochastic actor. In International conference on machine learning, pages 1861-1870. PMLR, 2018.
[20] Danijar Hafner, Timothy Lillicrap, Jimmy Ba, and Mohammad Norouzi. Dream to control: Learning behaviors by latent imagination. arXiv preprint arXiv: /ttp:///2..0000 2019.
[21] Danijar Hafner, Timothy Lillicrap. Ian Fischer, Ruben Villegas, David Ha, Honglak Lee, and James Davidson. Learning latent dynamics for planning from pixels. In International Confer- ence on Machine Learning, pages 2555-2565 PMLR, 2019.
[22] Song Han, Huizi Mao, and William J Dally. Deep compression: Compressing deep neural networks with pruning, trained quantization and huffman coding. arXiv preprint arXiv: 15 10,00149, 2015.

---

[23] Song Han, Jeff Pool, John Tran, and William Dally. Learning both weights and connections for efficient neural network. Advances in neural information processing systems, 28, 2015.
[24] Yihui He, Xiangyu Zhang, and Jian Sun. Channel pruning for accelerating very deep neural networks. In Proceedings of the IEEE international conference on computer vision, pages 1389-1397, 2017.
[25] Arthur J Helmicki, Clas A Jacobson, and Carl N Nett. Control oriented system identification: a worss-case/deterministt approach in h/sub infinity, IEEE Transactions on Automatic control, 16(10):1133-117, 1991.
[26] François Robert Hogan and Alberto Rodriguez. Feedback control of the pusher-slider system: story of hybrid and underactuated contact dynamics. arXiv preprint arXiv:161 2016.
[27] Zixuan Huang, Xingyu Lin, and David Held. Mesh-based dynamics model with occlusion reasoning for cloth manipulation. In Robotics: Science and Systems (RSS), 2022.
[28] Eric Jang, Shixiang Gu, and Ben Poole. Categorical reparameterizatioo with gumbel-softmax. arXiv preprint arXiv: /61 1.01 2016.
[29] Michael Janner, Justin Fu, Marvin Zhang, and Sergey Levine When to trust your model: Model-based policy optimization. CoRR, abs/1906.08253, 2019, URL http://arxiv..om org/ abs/1906. 08253.
[30] Natalie A Jones, Helen Ross, Timothy Lynam, Pascal Perez, and Anne Leitch. Mental models: an interdisciplinary synthesis of theory and methods. Ecology and Society, 16(1), 2011.
[31] Tejas D Kulkarni, Ankush Gupta, Catalin Ionescu, Sebastian Borgeaud, Malcolm Reynolds, Andrew Zisserman, and Volodymyr Mnih. Unsupervised learning of object keypoints for per- ception and control. Advances in neural information processing systems, 32:10724-10734, 2019.
[32] Brenden M Lake, Tomer D Ullman, Joshua B Tenenbaum, and Samuel J Gershman. Building machines that learn and think like people. Behavioral and brain sciences, 40, 2017.
[33] Yann LeCun, John S Denker, and Sara A Solla. Optimal brain damage. Advances in neural information processing systems, pages 598-605 1990.
[34] Joonho Lee, Jemin Hwangbo, Lorenz Wellhausen, Vladlen Koltun, and Marco Hutter. Learn- ing quadrupedal locomotion over challenging terrain. Science robotics, 5(47), 2020.
[35] Hao Li, Asim Kadav, Igor Durdanovic, Hanan Samet, and Hans Peter Graf. Pruning filters for efficient convnets. arXiv preprini arXiv:1600$ 08710, 2016
[36] Yunzhu Li, Jiajun Wu, Russ Tedrake, Joshua B Tenenbaum, and Antonio Torralba. Learning particle dynamics for manipulating rigid bodies, deformable objects, and fluids. arXiv preprint 1eX.:08810001666 2018
[37] Yunzhu Li, Jiajun Wu, Yan Zhu, Joshua B Tenenbaum, Antonio Torralba, and Russ Tedrake. Propagation networks for model-based control under partial observation. In 2019 International Conference on Robotics and Automation (ICRA), pages 1205-1211. IEEE, 2019.
[38] Yunzhu Li, Hao He, Jiajun Wu, Dina Katabi, and Antonio Torralba. Learning compositional koopman operators for model-based control. In International Conference Learning Repre- sentations, 2020. URL https: / //opereeiex/ wwwpeeeeeeeeeeeneeeeeeeeee Hi-1mmmmmddd111224PP
[39] Yunzhu Li, Antonio Torralba, Anima Anandkumar, Dieter Fox, and Animesh Garg. Causal dis- covery in physical systems from videos. Advances Neural Information Processing Systems, 33, 2020.
[40] Chenxi Liu, Barret Zoph, Maxim Neumann, Jonathon Shiens, Wei Hua, Li-Jia Li, Li Fei-Fei, Alan Yuille, Jonathan Huang, and Kevin Murphy. Progressive neural architecture search. In Proceedings of the European conference on computer vision (ECCV), pages 19-34, 2018.
12

---

[41] Hanxiao Liu, Karen Simonyan, and Yiming Yang. Darts: Differentiable architecture search. arXiv preprint arXiv:8006. 0900,, 2018.
[42] Zechun Liu, Haoyuan Mu, Xiangyu Zhang. Zichao Guo, Xin Yang, Kwang- Awang-Ting Cheng, and Jian Sun. Metapruning: Meta learning for automatic neural network channel pruning. In Proceedings of the IEEE/CVF international conference on computer vision, pages 3296-3305, 2019.
[43] Zhuang Liu, Jianguo Li, Zhiqiang Shen, Gao Huang, Shoumeng Yan, and Changshui Zhang. Learning efficient convolutional networks through network slimming. Proceedings of the IEEE international conference computer vision, pages 2736-2744, 2017.
[44]
Lennart Ljung. System identification. Springer, 1998.
[45] Miles Macklin, Matthias Müller, Nuttapong Chentanez, and Tae- Yong Kim. Unified particle physics for real-time applications. ACM Transactions on Graphics (TOG), 33(4): 2014.
[46] Chris J Maddison, Andriy Mnih, and Yee Whye Teh. The concrete distribution: A continuous relaxation of discrete random variables. arXiv preprint 1.00012,, 2016.
[47] Lucas Manuelli, Yunzhu Li, Pete Florence, and Russ Tedrake. Keypoints into the fu- ture: Self-supervised correspondence in model-based reinforcement learning. arXiv preprint MxXi:20099000885 2020.
[48] Tobia Marcucci and Russ Tedrake. Mixed-integer formulations for optimal control of piecewise-affine systems. In Proceedings of the 22nd ACM International Conference on Hy- brid Systems: Computation and Control, pages 230-239, 2019.
[49] Tobia Marcueci and Russ Tedrake. Warm start of mixed-integer programs for model predictive control of hybrid systems. IEEE Transactions on Automatic Control, 2020.
[50] Tobia Marcucci, Jack Umenberger, Pablo A Parrilo, and Russ Tedrake. Shortest paths in graphs of convex sets. arXiv preprint arXiv:2101. 2021.
[51] Microsoft. Neural Network Intelligence, 1 2021. URL / nni.
[52] Peter Mitrano, Dale McConachie, and Dmitry Berenson. Learning where trust unreliable models in an unstructured world for deformable object manipulation. Science Robotics, 6(54): eabd81 100, 2021.
[53] Thomas M Moerland, Joost Broekens, and Catholijn M Jonker. Model-based reinforcement learning: A survey. arXiv preprint vrXiv:2006.16711 2020.
[54] Pavlo Molchanov, Stephen Tyree, Tero Karras, Timo Aila, and Jan Kautz. Pruning convo- lutional neural networks for resource efficient inference. arXiv preprint arXiv: 1.06440, 2016.
[55] Bruce Moore. Principal component analysis in linear systems: Controllability, observability, and model reduction. IEEE transactions on automatic control, 26(1):1732,, 1981.
[56] Anusha Nagabandi, Kurt Konolige, Sergey Levine, and Vikash Kumar. Deep dynamics models for learning dexterous manipulation. In Conference on Robot Learning, pages 1101-1112. PMLR, 2020.
[57] Reuven Rubinstein and Dirk P Kroese. The cross-entropy method: a unified approach to combinatorial optimization, Monte-Carka simulation and machine learning. Springer Science & Business Media, 2013.
[58] Julian Schritwieser, Ioannis Antonoglou, Thomas Hubert, Karen Simonyan, Laurent Sifre, Si- mon Schmitt, Arthur Guez, Edward Lockhart, Demis Hassabis, Thore Graepel, et al. Mastering atari, go, chess and shogi by planning with a learned model. Nature, 388(7839)6004-69, 2020.
[59] John Schulman, Filip Wolski, Prafulla Dhariwal, Alec Radford, and Oleg Klimov. Proximal 2017.
policy optimization algorithms. arXiv preprint
13

---

[60] Haochen Shi, Huazhe Xu. Zhiao Huang, Yunzhu Li, and Jiajun Wu. Robocraft: Learn- ing to see, simulate, and shape clasto-plaste objects with graph networks. arXiv preprint arXiv:2205.02909, 2022.
[61] Kin Cheong Sou, Alexandre Megretski, and Luca Daniel. A quasi-convex optimization ap- proach 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: 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 Pro- ceedings 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: 2015.
[66] Wei Wen, Chunpeng Wu, Yandan Wang, Yiran Chen, and Hai Li. Learning structured sparsity 2016.
in deep neural networks. arXiv preprint
[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 Confer- ence on Computer Vision and Pattern Recognition, pages 10734-10744 2019.
[69] Philipp Wu, Alejandro Escontrela, Danijar Hafner, Pieter Abbeel, and Ken Goldberg. Day- dreamer: 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. Train- ing for faster adversarial robustness verification via inducing relu stability. arXiv preprint 1809::1009.03000 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 M fine-grad tinengrained structured sparse neural networks from scratch. arXiv preprint arXiv:2102.044
2021.
[73] Jiaji Zhou, Yifan Hou, and Matthew T Mason. Pushing revisited: Differential flatness, trajec- tory planning, and stabilization. The International Journal of Robotics Research, 38(12-13): 1477-1489, 2019.
[74] Barret Zoph and Quoe V Le. Neural architecture search with reinforcement learning. arXiv preprint 161100177 2016.
14

---

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 com- pare to prior works in model-based reinforcement learning, evaluated standard benchmark envi- ronments. We conduct experiments on two additional tasks from OpenAl Gym [7], CartPole-vl and Reacher-v4, following the based planner (MPPI) and a model-free RL method (PPO), we employ two additional model- -ased RL methods, (1) using PPO learn control policy from learned full neural dynamics model, and (2) MBPO [29] learning a model and a policy from scratch. The model-based RI. methods require additional time train policy using the learned dynamics model, whereas approach directly optimizes a task objective over the dynamics model without needing additional training.
procedures described in Section 4.3. On top of a sampling-
The experiment results shown in Figure 6 further demonstrate the superior performance of our ap- proach compared to prior methods on the two standard benchmark tasks. Notably, our approach achieves better performance with highly sparsified neural dynamics models with fewer ReLUs com- pared to prior works.
Reacher-v4
CartPole-v1
0.040.
0.25
MPP1
0.20
PPO
PPO+ Our Model
0.009
MPP1
PPO
RMSE
PPO Our Model MBPO
0.15
RMSE
0.007
0.10
0.005
0.05
0.00
0.003
162432
48
192
4
16
48
Remaining ReLUs
Remaining ReLUs
Figure 6: Closed-looy 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.
Closed-Loop Performance
0.125
0.100
RMSE
0.075
0.050
0.025
!
0.000
Remaining ReLUs
Figure 7: We tested the closed-loop control performance of dynamics models trained simplified using our method by incorporating them as the forward model in model-based RI. framework optimized with PPO. Our findings indicate that even when the dynamics models 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 generic and not limited to only working with our planning framework. We take the learned full and sparsified dynamics models the CartPole-vi environment and train control policy with PPO interacting only with the learned model, and provide the experiment results in Figure 7.
15

---

Ground Truth
Ours (2 ReLUs)
Li et all (2 UsL
Ground Truth
Our (2 ReLUs)
Li at al. (2 ReLUs)
(a) Piecewise affine function #1
(b) Piecewise affine function #2
Figure 8: Comparison with channel pruning baseline on recovering 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 basclinc does not provide an accurate match.
The results show that the neural dynamics models trained in 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 sparsífication technique with the pruning meth- ods 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. con- trast, our method aims eliminate as many nonlinearities from the neural networks possible. This differs from channel pruning, which only zeroes out values. Our approach permits the replace- ment of ReLU activations with identity mappings, the inclusion of which allows more accurate model to be achieved at an equivalent level of sparsification. This offers considerable advantage during the planning stage.
To illustrate our point more concretely, we provide, in this section, experimental results comparing method against Level Channel Pruning referenced in [35].
C.1 Evaluation on Piecewise Affine (PWA) Functions
The two pruning methods are tasked recover the ground truth PWA functions. detailed in the experiment section of the main paper. Figure 8 illustrates the results after sparsifying the neural networks to two rectified linear units (ReL.U) using both methods. Our method successfully identi- fies 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, train the neural dynamics model for 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] 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 similar procedure to construct the baseline. During the initial model training phase, the mask resulting from the Lı 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 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 ReL. units.
16

---

Object Pushing
Object Sorting-2
Object Sorting-3
Rope Manipulation
0.5
Ours
0.40
Ours
0.6
Ours
Ours
et.al
ct al.
et
Li
al
0.30
0.6
0.25
0.20
0.4
0.15
0.05
0.1-
0.0
0.00
0.0
12
24
36
60
100
12
24
28
100
6
12
24 48
100
12
24
36
60 100
Remaining ReLUs
Remaining ReLUs
Remaining ReLUs
Remaining ReLUs
Figure 9: Dynamics prediction error 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 state. Because 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 sparsification method and the baseline in Fig- ure 9. Throughout the sparsification process, because our sparsífication objective allows replacing non-linearities with identity mappings, our method consistently achieves superior performance measured by prediction error, across all tasks.
C.3 Evaluation on Closed-Loop Control
In this experiment, aim further examine whether method also boosts the closed-loop control performance that critical for executing the optimized plans in the real world. We choose the object pushing task, and prune the learned dynamics model down 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 method consistently exhibit superior closed-loop control performance across all sparsification levels.
Object Pushing
0.6
Ours
Li et al.
0.5
RMSE
0.4
0.2
0.0
15
18
24
36
Remaining ReLUs
Figure 10: An ablation of sparsification method compared with a prior network pruning method, evaluated by closed-lowe loop control performance, demonstrating superior performance closed-loop control.
17


Reviews:
[{"review_id": "asWBrb8O3Q", "reviewer": "Reviewer_h4y2", "summary": "This paper focused on the combination of predictive control and model learning. An autogressive dynamic model based on a ReLU neural network is first learned over the observation space. The authors then aimed to sparsify it after introducing the indicator mapping function. To make the optimization feasible, the Gumbel-Softmax trick is applied to replace the greedy operation. Once the sparse dynamic is learned, mixed-integer programming solvers are used to obtain the control policy. Finally, the authors validated the proposed method across a few tasks and showed its promising performance.", "questions": "1. What's the complexity of applying MIP?\n\n2. What's the underlying architecture for the ReLu neural network? It would be a bit surprising if a feedforward neural network can capture the features from image inputs.", "limitations": "Yes", "rating": 5, "confidence": 3, "soundness": 3, "presentation": 2, "contribution": 2, "strengths": "1. The proposed framework of sparse neural dynamics in predictive control looks interesting. The use of Gumbel-Softmax to reparameterizing the original discrte optimization makes the gradient based methods feasible, and is indeed a reasonable idea. \n\n2. The optimization for the control part is more efficient, compared with the gradient based methods, which constitutes another contribution of this paper. \n", "weaknesses": "Regarding the experimental comparison, it would be more convincing if the authors can test on commonly used reinforcement learning benchmarks. The only comparison regarding control policy is from Figure 3 (b) and it only involves different optimization solvers. The current results would leave the impression that the proposed method may overfit on these tasks only.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time reviewing our paper, and for your insightful suggestions that have improved our work.\n> It would be more convincing if the authors can test on commonly used reinforcement learning benchmarks\n\nWe conducted further experiments on two additional environments, Reacher-v4 and Cartpole-v1, from OpenAI Gym [3], and reported the results below (also in Rebuttal PDF Fig. 1)\n|                  |             Ours            |        MPPI       | Model-free PPO | PPO + Our model |     MBPO [4]    |\n|-------------|:-----------------------:|:---------------:|:--------------:|:---------------:|:---------------:|\n| Cartpole-v1 |  **0.003728** (4 ReLUs) | 0.003735 (full) |    0.039222    | 0.006291 (full) | 0.004268 (full) |\n|  Reacher-v4 | **0.053467** (48 ReLUs) | 0.064792 (full) |    0.237225    | 0.108016 (full) |  0.110031(full) |\n\nEach number represents the median of RMSE over 100 trials. Parentheses denote either the method used a full model or a sparsified model with specified number of ReLUs.\n\nAs the table illustrates, our approach outperforms prior methods on the two standard RL benchmark environments. Notably, our approach achieved superior performance with highly sparsified neural dynamics models with fewer ReLUs compared to prior works.\n\n> The current results would leave the impression that the proposed method may overfit on these tasks only.\n\nOur experiments covered a variety of tasks, involving rigid and deformable objects in scenarios ranging from single to multi-objects, applied to dynamics models instantiated using feed-forward neural networks and graph-neural networks. These experiments demonstrated the applicability and combinatorial generalizability of our method. We hope that the results on the two additional tasks further substantiate that our method is applicable to a wide range of task settings.\n\n> What's the complexity of applying MIP?\n\nSolving MIPs is NP-hard, so in the worst case, the solve time might be exponential in the number of ReLUs. However, there exist highly optimized solvers like Gurobi that can solve medium-size MIPs very quickly in practice. The solution of a MIP can also be terminated early, and the branch and bound solver will return a solution together with an upper bound on the distance of this solution from the global minimum. This strategy can be leveraged to further reduce the runtimes, while still providing strong optimality guarantees. We provide details on various strategies in the literature to accelerate solving MIPs in Section 3.5.1 of the paper.\n\n> What's the underlying architecture for the ReLu neural network? It would be a bit surprising if a feedforward neural network can capture the features from image inputs.\n\nBelow are the architecture details (also provided in supplemental material, section B). The ordered lists represent the number of ReLUs in each layer of the network.\n\n- PWA functions: 96, 192, 192, 96\n\n- Object pushing: 256, 256, 256\n\n- Object sorting: graph neural network similar to Sanchez-Gonzales et al. [5] with 64 hidden units per layer, with a total of 512 ReLUs\n\n- Rope manipulation: 256, 256, 256\n\n- Cartpole-v1: 16, 16, 16\n\n- Reacher-v4: 64, 64, 64\n\nWe used either keypoint (object pushing, rope manipulation), object-centric representation (object sorting), or simulation state (cartpole, reacher) as our state representations. Prior works [1, 2] demonstrated that low-dimensional structured representations have superior generalization capabilities compared to pixel-based dynamics, especially for compositional systems containing varying numbers of objects (as is the case in our Object Sorting tasks). \n\n**Have these responses addressed the reviewer’s concerns? We look forward to continuing the discussion.**\n\n[1] N. Watters, A. Tacchetti, T. Weber, R. Pascanu, P. W. Battaglia, and D. Zoran, ‘Visual Interaction Networks’, CoRR, vol. abs/1706.01433, 2017.\n\n[2] D. Driess, Z. Huang, Y. Li, R. Tedrake, and M. Toussaint, ‘Learning Multi-Object Dynamics with Compositional Neural Radiance Fields’, in Conference on Robot Learning, 2022.\n\n[3] G. Brockman et al., ‘OpenAI Gym’, CoRR, vol. abs/1606.01540, 2016.\n\n​​[4] M. Janner, J. Fu, M. Zhang, and S. Levine, ‘When to Trust Your Model: Model-Based Policy Optimization’, CoRR, vol. abs/1906.08253, 2019.\n\n[5] A. Sanchez-Gonzalez et al., ‘Graph networks as learnable physics engines for inference and control’, CoRR, vol. abs/1806.01242, 2018."}, {"review_id": "B9PWHBvtUQ", "reviewer": "Reviewer_H7kE", "summary": "This paper proposes a framework for model-based planning with forward dynamics represented as sparse neural networks. The paper examines different ways of inducing sparsity in MLP and GNN based forward models, and performs real robot manipulation experiments investigating the tradeoffs with sparsity and performance. ", "questions": "Please refer to my list of weaknesses above. In summary\n\n- comparisons to prior MBRL approaches?\n- discussion about generality of the findings?\n- feasibility of experiments with image observations?", "limitations": "The limitations are a bit unclear at the moment (also refer to my list of weaknesses). It will be helpful to provide a detailed discussion of the limitations and generality of the findings. ", "rating": 6, "confidence": 4, "soundness": 3, "presentation": 3, "contribution": 3, "strengths": "- the paper targets an interesting problem of learning efficient world models for control, by reducing optimization costs in learning the forward model. The specific approach of investigating how sparsity can enable this efficient learning of the forward model for real robot manipulation scenarios, is novel in my understanding.\n\n- the paper is easy to follow, with detailed descriptions of the different ways of inducing sparsity, and the different architectures. The research questions investigated are well formulated, and adequately addressed in the experiments\n\n- the experiments are on interesting real robot manipulation tasks like object sorting, and rope manipulation, so the findings are amenable to be deployed in real world control scenarios. ", "weaknesses": "- the main weakness of the paper is that there aren't comparisons to prior model-based RL approaches, e.g. PETS, Dreamer, MBPO etc (https://arxiv.org/abs/1805.12114,https://arxiv.org/abs/1912.01603, https://arxiv.org/abs/1906.08253)  . If it is difficult to compare them on hardware experiments,  there should be comparisons done on simulated robot control tasks. \n\n- it is unclear how general are the findings in terms of being applied to different model-based control frameworks. There should be a disucssion (and possibly experiments) about if we can take any prior model-based RL algorithm and replace the forward model with a sparse neural network, and see benefits under certain assumptions - if not, then a discussion of what exactly are the requirements on how the model should be learning and how a policy (or plan) should be learned is needed.\n\n- the experiments results in my understanding are on low-dimensional states instead of images. Is there a reason for not evaluating with image observations? Encoding image observations into latent states, and then doing policy learning is a common practice in the community now - so it will be helpful to have a discussion regarding this.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper, and for your thoughtful comments and suggestions that have contributed to the refinement of our work.\n> there aren't comparisons to prior model-based RL approaches, e.g. PETS, Dreamer, MBPO etc\n\nWe conducted additional experiments employing two model-based RL methods, 1) our learned full neural dynamics model with PPO to learn a control policy, and 2) MBPO [4] learning a model and a policy from scratch, on two OpenAI Gym [3] benchmark environments, Cartpole-v1 and Reacher-v4. Below are the results (also included in Rebuttal PDF Fig. 1) \n\n|             |           Ours          |       MPPI      | Model-free PPO | PPO + Our model |     MBPO [4]    |\n|-------------|:-----------------------:|:---------------:|:--------------:|:---------------:|:---------------:|\n| Cartpole-v1 |  **0.003728** (4 ReLUs) | 0.003735 (full) |    0.039222    | 0.006291 (full) | 0.004268 (full) |\n|  Reacher-v4 | **0.053467** (48 ReLUs) | 0.064792 (full) |    0.237225    | 0.108016 (full) |  0.110031(full) |\n\nEach number represents the median of RMSE over 100 trials. The parenthesis after the number denotes the number of ReLUs used in the neural dynamics model.\n\nThe 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.\n\nOverall, our experiment results demonstrate that our method using sparsified neural dynamics models with fewer ReLUs can be applied to a wide variety of tasks and exhibit better performance compared to prior methods examined.\n\n> it is unclear how general are the findings in terms of being applied to different model-based control frameworks\n\nThe neural dynamics model learned in our method is generic and not limited to only working with our planning framework. We took the learned full and sparsified dynamics models and trained a control policy with PPO interacting only with the learned model as suggested, and provide the experiment results below (also included in Rebuttal PDF Fig. 2).\n\n| Num. ReLUs |    48    |    16    |     8    |     4    |     2    |     1    |     0    |\n|:----------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|\n|    RMSE    | 0.006291 | 0.007201 | 0.009670 | 0.018122 | 0.088279 | 0.088655 | 0.089807 |\n\nThe results above on Cartpole-v1 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.\n\n> Is there a reason for not evaluating with image observations?\n\nOur method is agnostic to model architectures and can be applied to incorporate image observations encoded as latent states as the reviewer suggested, as long as the encoder is a ReLU network. For the specific real-world robotic manipulation tasks we considered in this paper, prior works [1, 2] showed that using low-dimensional structured representations provides stronger generalization capabilities compared to pixel-based dynamics, especially for compositional systems containing varying numbers of objects (as is the case in our Object Sorting tasks). That’s why we chose to not evaluate with image observations in our experiments.\n\n> It will be helpful to provide a detailed discussion of the limitations and generality of the findings.\n\nPlease refer to the global rebuttal response for discussions of limitations and generality of our work.\n\n**Have these responses addressed the reviewer’s concerns? We look forward to continuing the discussion.**\n\n[1] N. Watters, A. Tacchetti, T. Weber, R. Pascanu, P. W. Battaglia, and D. Zoran, ‘Visual Interaction Networks’, CoRR, vol. abs/1706.01433, 2017.\n\n\n[2] D. Driess, Z. Huang, Y. Li, R. Tedrake, and M. Toussaint, ‘Learning Multi-Object Dynamics with Compositional Neural Radiance Fields’, in Conference on Robot Learning, 2022.\n\n\n[3] G. Brockman et al., ‘OpenAI Gym’, CoRR, vol. abs/1606.01540, 2016.\n\n\n​​[4] M. Janner, J. Fu, M. Zhang, and S. Levine, ‘When to Trust Your Model: Model-Based Policy Optimization’, CoRR, vol. abs/1906.08253, 2019.\n"}, {"review_id": "l6veqm4cWp", "reviewer": "Reviewer_fLUK", "summary": "This paper proposes a new framework for model-based control. The approach focuses on learning a sparse deep neural network and using a mixed-integer program solver for closed-loop planning. Experimental results are presented on several tasks including object and rope manipulation tasks. The results show that the proposed approach improves performance over strong baselines.", "questions": "I am willing to adjust my rating if these questions can be answered:\n\n- Please comment further on the results in Figure 5; there is a trend but why is this a significant advance?\n- Does the proposed sparsification technique improve closed-loop performance?\n- What are the limitations of the work?", "limitations": "No, the paper does not discuss limitations.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 4, "contribution": 4, "strengths": "- The framework proposed in this paper elegantly combines concepts from deep neural network (DNN) pruning and mixed integer programming (MIP) into a solution for model-based control.\n- The specific technique for sparsifying a DNN by removing non-linearities rather than simply dropping nodes intuitively matches the goals of eventually using the model with a MIP solver.\n- The experiments demonstrate that the approach is effective in both simulation and in the real-world on real hardware -- providing strong evidence that this is a generally applicable approach.\n- The technique is demonstrated with both MLP and GNN based models, showing the versatility of the proposed approach.\n- The supplemental video presents strong qualitative evidence supporting the efficacy of the approach on real-robots.\n- The paper is very well written.", "weaknesses": "- It is unclear if the performance gains in Figure 5 are significant. On the Object Pushing and Rope Manipulation tasks MIP does not appear to outperform MPPI. And the improvements of MIP on the Object Sorting tasks appear to be within the error bars.\n- A key ablation is missing. One of the main claims in this work is that the proposed sparsification technique, which focuses on removing nonlinearities rather than neurons, improves performance. Only a partial ablation in support of this claim is provided in the appendix A2. Specifically, A2 shows that the proposed approach leads to lower prediction error. However, to quote L277-279 in the main paper, “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.” I agree. Thus, it is similarly crucial to perform this ablation in a closed-loop setting.", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for reviewing our paper, and for your insightful feedback that helped improve our work.\n\n> Please comment further on the results in Figure 5; there is a trend but why is this a significant advance?\n\nIn both Object Sorting tasks, our method using MIP and a sparsified model of 24 and 36 ReLUs demonstrate superior performance compared to the method using MPPI and a full model with 512 ReLUs, indicated by the lower interquartile range and median, supporting our statement that having a sparsified model allows us to benefit from more powerful optimization tools, leading to superior closed-loop control results despite having a less accurate dynamics model. \n\nIn Object Pushing, MIP with a model of 60 ReLUs achieved comparable performance to MPPI with a model of 768 ReLUs, and slightly better performance compared to MPPI with 200 ReLUs. \n\nWe agree that MIP achieves a similar performance to MPPI on Rope Manipulation. There is a more significant gap in model prediction error between the full dynamics model and highly sparsified models due to the complex physical properties of the deformable rope (Fig. 3a). The task requires accurately manipulating the rope to match a target shape specified with key points along the rope, thus the performance is highly dependent on having an accurate dynamics model. Even though using a sparsified model enabled us to leverage a more powerful optimization tool giving better results in open-loop optimization (Fig. 3b), the model prediction error limits the performance upper bound in closed-loop planning. Future improvements could involve the co-optimization of the model sparsification procedure with control synthesis. Enabling the sparsification process to be mindful of the subsequent control task may lead to even better performance\n\n> Does the proposed sparsification technique improve closed-loop performance?\n\nThe reviewer pointed out that it is crucial to perform an additional ablation study of our proposed sparsification technique in a closed-loop setting, as we also stated in the paper. To that end, we performed additional experiments to validate the performance again on closed-loop control of our proposed method, and provide the results below (also in Rebuttal PDF Fig. 3 right).\n\n| Num. ReLUs |    36    |    24    |    18    |    15    |\n|:----------:|:--------:|:--------:|:--------:|:--------:|\n|    Ours    | 0.172873 | 0.192272 | 0.241469 | 0.271934 |\n|  Li et al. [1] | 0.186144 | 0.239412 | 0.321603 | 0.287213 |\n\nThe numbers in each column correspond to the closed-loop control performance, measured by RMSE, using the sparsified model with the corresponding number of ReLUs.\n\nThese ablation results, combined with results reported in supplemental material Fig. 2, demonstrate that our proposed sparsification scheme brings significant improvements in both open-loop prediction and closed-loop control performance.\n\n> What are the limitations of the work?\n\nPlease refer to the global rebuttal response for discussions of limitations of our work.\n\n**Have these responses addressed the reviewer’s concerns? We look forward to continuing the discussion.**\n\n[1] H. Li, A. Kadav, I. Durdanovic, H. Samet, and H. P. Graf, ‘Pruning Filters for Efficient ConvNets’, CoRR, vol. abs/1608.08710, 2016."}, {"review_id": "aTQLvb8vdf", "reviewer": "Reviewer_LeuW", "summary": "This paper proposes a method for pruning neural networks with ReLu activation functions during training. When employed for learning the dynamics of a control system, this often leads to networks with few activation functions performing similarly as large networks. This allows to apply mixed integer programming techniques to determine optimal control policies. The effectiveness of this approach is demonstrated in a comparison with a sampling based optimizer on several tasks for a robotic manipulator.", "questions": "What is the point of the regularization in (9) if you bound the total number of ReLu activation functions through $\\epsilon$ in (5) anyways?\n\nWhat is a closed-loop planning performance? Or do you mean closed-loop control performance? \n\nIn (8), is $z_{ij}$ the approximation of $\\pi_{ij}$?", "limitations": "Potential negative societal impact is not mentioned in the paper. Some limitations of the proposed approach (e.g., restriction to ReLu activation functions) are mentioned throughout the paper, but there is no dedicated paragraph about limitations. I think there are some limitations (e.g., computational complexity) that should be added.", "rating": 5, "confidence": 3, "soundness": 2, "presentation": 3, "contribution": 2, "strengths": "The idea of formulating planning problems with neural network dynamics as mixed integer program seems novel to me. Moreover, I find the proposed approach of reducing the number of neurons during training very interesting, but I cannot comment on its novelty since I am no expert on this topic. Overall, the paper reads very well and has a good structure. The nice demonstration of the proposed method in a real-world robotic experiment is also a strength of the paper.", "weaknesses": "My biggest concern is the missing discussion of computational complexity/computation time. I am not an expert on mixed integer programming, but a quick search suggests that even mixed integer linear programming is already NP hard, so this seems to be a problem in general. I understand that this problem can be mitigated through a sufficiently small number of neurons (‘If only a few ReLUs are left in the model, Equation 10 can be efficiently solved to global optimality’), but I am missing a clear specification what sufficient means in this context. Moreover, I would expect it to be connected to the prediction horizon, i.e., the number of time steps considered in the sum in (10). Therefore, I think a complexity (e.g., in O notation) should be provided to give the reader an impression how severe the computation times grows. This weakness also limits the usefulness of the evaluation in my opinion. Why would you reduce the number of neurons for MPPI in practice? I think the performance in relation to computation time is rather the important metric to look at for this comparison, i.e., a complexity-performance trade-off comparison would be essential. Therefore, the comparison seems a little unfair at the moment. This similarly extends to the robot experiment in Sec. 4.3, where the sampling and the horizon rate of the MPC are not specified. Moreover, it is not clear how fast the robot moves. These are all parameters which crucially influence how challenging the problem is.\n\nIn general, the method seems tailored to piece-wise affine systems or systems that look almost like that. The examples seem to go into this direction, but I cannot say for sure since I could not find information what dynamics are actually learned in the robotic examples. I think it would be interesting to see how many activation functions are needed for accurately learning highly nonlinear dynamics, e.g., cart-pole swing-up, and how this affects the control performance. \n\nOverall, the novelty seems to be mainly the (straightforward) connection of existing ideas, even though I admit that it is a very clever combination. The proposed reduction of neurons is not even targeted to achieve a high control performance, but only to maintain high model accuracy. High control performance is only addressed a posteori by selecting the best model observed in experiments. Doing it as described in lines 206-208 runs the risk of executing a potentially bad controller on a real system. This seems like a dangerous thing to generally do. I think it would be much more interesting to directly optimize the model for achieving the best control performance.\n\nFinally, I find the comparison in the numerical evaluation a little weak. When no experiment has to be done like for the open-loop planning performance evaluation, I think more methods than only MPPI should be investigated. I am not an expert, but I know there exists more than one planning/control method for nonlinear systems, e.g., MPC with NN dynamics (Salzmann et al., 2023). Moreover, I do not understand why the learned dynamics model is not used as environment for the classical RL methods. It is apparently accurate enough to allow for a direct transfer to a real-world experiment. So why are the interactions limited to a relatively small number with the real environment, when a large number of interactions with a potentially only slightly more inaccurate model are available?\n\nT. Salzmann, E. Kaufmann, J. Arrizabalaga, M. Pavone, D. Scaramuzza and M. Ryll, \"Real-Time Neural MPC: Deep Learning Model Predictive Control for Quadrotors and Agile Robotic Platforms,\" in IEEE Robotics and Automation Letters, vol. 8, no. 4, pp. 2397-2404, April 2023", "quality": null, "clarity": null, "significance": null, "originality": null, "strengths_and_weaknesses": "", "was_revised": false, "final_justification": "", "rebuttal": "Thank you for your time reviewing our paper, and for your constructive suggestions that helped improve our work.\n\n> I think a complexity (e.g., in O notation) should be provided to give the reader an impression how severe the computation times grows.\n\nWe agree with the reviewer that a more precise description of the runtime is useful for the readers. The complexity is exponential in the number of ReLUs in the worst case, but in practice solvers like Gurobi are highly optimized, and their branch-and-bound implementations can solve medium-scale problems very quickly, although the performance highly depends on the problem formulation, problem conditioning, and the tightness of the convex relaxation.\n\n> I think the performance in relation to computation time is rather the important metric to look at for this comparison, i.e., a complexity-performance trade-off comparison would be essential.\n\nWe agree that a complexity-performance trade-off is essential, and we used the number of remaining ReLUs in the dynamics model as a proxy for computation complexity. The solve time in practice depends on the problem formulation, conditioning, initialization, and various solver techniques and heuristics. In our experiments, solving a MIP with Gurobi induces a high variance in computation time across different tasks, dynamics model architectures, and problem initializations, as shown in Rebuttal PDF Fig. 3 left. \n\nAlthough the wall clock computation time exhibits high variance for the same neural dynamics model, we observe that the computation time is positively correlated with the number of ReLUs in the network, and thus chose to use the number of remaining ReLUs as an effective proxy for computation complexity in our experiments.\n\n> the method seems tailored to piece-wise affine systems or systems that look almost like that\n\nOur method focuses on neural networks with ReLU activations, which act as piecewise affine functions. With a sufficient number of pieces, they can exhibit extremely strong approximation power and can be used to approximate nonlinear functions that are smooth arbitrarily well. Piecewise affine systems have been widely studied and used in the literature [Sontag et al.] to approximate highly complex and nonlinear systems. \n\nOur experiments examined the applicability of our approach to systems that are not piecewise affine. The dynamics of the rope in our Rope Manipulation task are not piecewise affine. The Object Pushing task involves modeling the object’s orientation, which is also not piecewise affine.\n\n> The proposed reduction of neurons is not even targeted to achieve a high control performance, but only to maintain high model accuracy.\n\nUsing the same control optimization tool, we believe that a dynamics model with higher accuracy leads to better control performance as long as the model size does not need to be too large for the control optimization tool to solve within reasonable time. Our experiments shown in Fig 5 and Rebuttal PDF Fig. 1 also show that the closed-loop control performance decreases with the model accuracy, using the same optimization tool. \n\nWe would also like to note that targeting to simultaneously achieve a high control performance during sparsification might necessitate a differentiable design for the control optimization procedure, which we will leave for future work.\n\n> I think more methods than only MPPI should be investigated\n\nWe provided comparisons to MPPI, as well as model-free RL methods (PPO and SAC), all of which are widely used in the literature solving similar control and manipulation tasks. We also performed additional experiments using model-based RL (MBPO [Janner et al.]) on two OpenAI Gym benchmark environments (Cartpole-v1, Reacher-v4). The results are shown in Rebuttal PDF Fig. 1. \nOur approach outperforms prior methods on the two standard RL benchmark environments. Notably, our approach achieved superior performance with highly sparsified neural dynamics models with fewer ReLUs compared to prior works.\n\n> I do not understand why the learned dynamics model is not used as environment for the classical RL methods.\n\nThe neural dynamics model learned in our method is generic and not limited to only working with our planning framework. We took the learned full and sparsified dynamics models and trained a control policy with PPO interacting only with the learned dynamics model as suggested, and provide the experiment results below (also included in Rebuttal PDF Fig. 2).\n\n| Num. ReLUs |    48    |    16    |     8    |     4    |     2    |     1    |     0    |\n|:----------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|:--------:|\n|    RMSE    | 0.006291 | 0.007201 | 0.009670 | 0.018122 | 0.088279 | 0.088655 | 0.089807 |\n\nThe results above on Cartpole-v1 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. Our best performing model, MIP with 4 ReLUs, outperformed these models and achieved an RMSE of 0.003728.\n\n> why are the interactions limited to a relatively small number with the real environment, when a large number of interactions with a potentially only slightly more inaccurate model are available?\n\nWe limited the number of interactions for training with the two model-free methods, PPO and SAC, to be the same as the number of real environment interactions that our method used to train the dynamics model. This setting is practical in scenarios where interacting with the real environment is costly or not always available, and is widely studied in the RL community. As the reviewer suggested, we conducted additional experiments using a model-based RL method to interact with the neural dynamics model to learn a control policy, reported above. \n \n**Due to length limit, we are happy to respond to the remaining questions during the discussion phase.**"}]

Author rebuttal:
We thank the reviewers for dedicating their time and effort in reviewing our paper, and we deeply appreciate their thoughtful comments and insightful feedback. We appreciate the reviewers agreeing that our approach is novel, well formulated, tackles an interesting problem, and that our paper is well written. 

> **Additional Experiments**

We have conducted additional experiments as requested by the reviewers and have included the results in the Rebuttal PDF, including:

1. Experiments involving two environments from OpenAI Gym (Cartpole-v1, Reacher-v4), to further demonstrate the generalizability of our method to a variety of tasks. (Reviewer h4y2)
2. Comparisons to an additional prior model-based RL method (MBPO), to showcase the superior performance of our method using a sparsified dynamics model compared to prior methods using a full dynamics model. (Reviewers H7kE, LeuW)
3. Taking dynamics models trained and sparsified through our method to train a control policy with PPO, demonstrating that our sparsification scheme produces dynamics models that are general and can be effectively combined with prior model-based RL methods. (Reviewers H7kE, LeuW)
4. Closed-loop control performance of dynamics models sparsified using our approach compared to models sparsified using a prior method, showing that our sparsification approach generates dynamics models with lower prediction errors and better closed-loop control performance. (Reviewer fLUK)
5. Experiments examining the positive correlation between computation complexity and the number of remaining ReLUs in the dynamics model, to further support using remaining ReLUs as an effective proxy for computation complexity. (Reviewers h4y2, LeuW)

> **Applicability of the Proposed Method**

Following the comments from Reviewers h4y2, H7kE, and LeuW about the applicability of our method, we hope to highlight that our method is generic and applicable to a wide variety of dynamics models instantiated using ReLU networks including but not limited to feed-forward neural networks and graph-neural networks examined in this paper. Our experiments also showed competitive performance and combinatorial generalizability on a wide range of tasks involving rigid and deformable objects in scenarios ranging from single to multi-objects.


> **Limitations**

In response to reviewers’ requests, we provide more details on the limitations of our approach below. We will include these discussions in the revised version of our paper.

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. We provided a review of a variety of strategies in the literature to accelerate the process of finding a solution for our MIPs in section 3.5.1 of the paper.

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. 

In terms of generality, our method naturally applies to dynamics models represented by ReLU networks with diverse model architectures (MLP, GNN, etc.), and showed superior performance on a variety of tasks involving rigid and deformable objects in single to multi-object scenarios, demonstrating the applicability and combinatorial generalizability of our method.

**We look forward to discussing further with the reviewers during the discussion phase.**

Meta-review: Thank you for your submission and active engagement throughout the review period. After the discussion period, the reviewers and I are in agreement that the paper will be a great addition to the conference. The idea of learning/pruning for sparse dynamics that can be used with a MIP planner is appealing and the experiments in comparison to MPPI and other baselines are well-executed. Two concerns came up during the discussion period that we encourage the authors to incorporate into the final version of the paper:

1.  On the ability of obtaining a globally optimal solution, the paper states "if only a few ReLUs are left in the model, Equation 10 can be efficiently solved to global optimality." However, depending on the cost function, even if the dynamics are formed as MIP, global optimality is not easy. Restricting the model flexibility to enable efficient planning is not new. For example, one can also exploit deep Koopman operator [1], whose planning procedure can be completed via LP, instead of MIP in this paper. However, this has not been discussed.
2. While the experimental demonstrations on the robotic environments are compelling, the fact that they are non-standard makes it difficult to understand if the method would work in other control and model-based reinforcement learning settings where the dynamics models are learned.

[1] Learning Compositional Koopman Operators for Model-Based Control. Yunzhu Li, Hao He, Jiajun Wu, Dina Katabi, Antonio Torralba
Decision: Accept (poster)