SlowGuess/NeurIPS_2023_OCR huggingface dataset:
Ziang Liu $^{1,2}$  Genggeng Zhou $^{2*}$  Jeff He $^{2*}$  Tobia Marcucci $^{3}$   
Li Fei-Fei $^{2}$  Jiajun Wu $^{2}$  Yunzhu Li $^{2,4}$ $^{1}$ Cornell University  ${}^{2}$ Stanford University  
 $^{3}$ Massachusetts Institute of Technology  
 $^{4}$ University of Illinois Urbana-Champaign  
ziangliu@cs.cornell.edu  
{g9zhou,jeff2024}@stanford.edu  
tobiam@mit.edu  
{feifeili, jiajunwu}@cs.stanford.edu  
yunzhuli@illinois.edu
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.
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]. It is, therefore, desirable to build predictive models of the environment from observations and develop 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 is hard and sometimes impossible to obtain in complicated robotic (manipulation) tasks. Learning-based dynamics modeling circumvents the problem by learning a predictive model directly from raw sensory observations, 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,
$^{\dagger}$  Please see our website at robopol.github.io/Sparse-Dynamics/ for additional visualizations.
(a) Sparsify the learned neural dynamics
(b) Planning with sparsified neural dynamics
Piecewise Affine Systems
Object Pushing
Object Sorting
(c) Benchmark tasks
Figure 1: Model-based control with sparse neural dynamics. (a) Our framework sparsifies the neural dynamics models by either removing neurons or replacing ReLU activation functions with identity mappings (ID). (b) The sparsified models enable 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 a deformable rope.
Rope Manipulation
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 better guarantees, but often struggle to accurately fit observation data. An important question arises: how precise do these models need to be when employed within a feedback control loop? The cognitive 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 environmental feedback [30, 10]. This notion is also key in control-oriented system identification [25, 44] and model order reduction [55, 61]. The framework from Li et al. [38] trades model expressiveness and precision for more efficient and effective optimization-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 1a). This yields a highly sparsified neural dynamics model, that is amenable to model-based control using state-of-the-art 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 piecewise affine systems. Moreover, it 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 be summarized as follows: (i) We propose a novel formulation for identifying the dynamics model from observation data. For this step, we introduce a continuous approximation of the sparsification problem, enabling end-to-end gradient-based optimization for both the model class and the model parameters (Figure 1a). (ii) By having significantly fewer ReLU units than
the full model, the sparsified dynamics model allows us to solve the predictive-control problems using efficient MIP solvers (Figure 1b). This can lead to better closed-loop performance compared to both model-free and model-based reinforcement learning (RL) 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 a variety of simulated and real-world manipulation tasks with complex contact dynamics, such as object pushing and sorting, and manipulation of deformable objects (Figure 1c).
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]. Recent 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, 39], 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 a 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, our 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 a neural network with ReLU activations is a piecewise affine function that can be modeled exactly through a set of mixed-integer linear inequalities. This allows us to use highly-effective 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].
In this section, we describe our methods for learning a dynamics model using environmental observations and for sparsifying DNNs through a 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  $\mathcal{D} = \{(y_t^m, u_t^m) \mid t = 1, \dots, T, m = 1, \dots, M\}$  collected via interactions with the environment, where  $y_t^m$  and  $u_t^m$  denote the observation and action obtained at time  $t$  in trajectory  $m$ . Our goal is to learn an autoregressive model  $\hat{f}_{\theta}$ , parameterized by  $\theta$ , as a proxy of the real dynamics that takes a small sequence of observations and actions from time  $t'$  to the current time  $t$ , and predicts the next observation at time  $t + 1$ :
$$
\hat {y} _ {t + 1} ^ {m} = \hat {f} _ {\theta} \left(y _ {t ^ {\prime}: t} ^ {m}, u _ {t ^ {\prime}: t} ^ {m}\right). \tag {1}
$$
We optimize the parameter  $\theta$  to minimize the simulation error that describes the long-term discrepancy between the prediction and the actual observation:
$$
\mathcal {L} (\theta) = \sum_ {m = 1} ^ {M} \sum_ {t} \| y _ {t + 1} ^ {m} - \hat {f} _ {\theta} \left(\hat {y} _ {t ^ {\prime}: t} ^ {m}, u _ {t ^ {\prime}: t} ^ {m}\right) \| _ {2} ^ {2}. \tag {2}
$$
3.2 Neural network sparsification by removing or replacing ReLU activations
We instantiate the transition function  $\hat{f}_{\theta}$  as a ReLU neural network with  $N$  hidden layers. Let us denote the number of neurons in the  $i^{\mathrm{th}}$  layer as  $N_{i}$ . When given an input  $x = (y_{t:t}^{m}, u_{t:t}^{m})$ , we denote the value of the  $j^{\mathrm{th}}$  neuron in layer  $i$  before the ReLU activation as  $x_{ij}$ . Regular ReLU neural networks apply the rectifier function to every  $x_{ij}$  and obtain the activation value using  $x_{ij}^{+} = h_{ij}(x_{ij}) \equiv \operatorname{ReLU}(x_{ij}) \triangleq \max(0, x_{ij})$ . 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_{ij}$  with either an identity mapping  $\operatorname{ID}(x_{ij}) \triangleq x_{ij}$  or a zero function  $\operatorname{Zero}(x_{ij}) \triangleq 0$ , where the latter is equivalent to removing the neuron (Figure 1a).
We divide the parameters in  $\hat{f}_{\theta}$  into two vectors,  $\theta = (\omega, \alpha)$ . The vector  $\omega$  collects the weight matrices and the bias terms. The vector  $\alpha$  consists of a set of integer variables that parameterize the architecture of the neural network:  $\alpha = \{\alpha_{ij} \in \{1, 2, 3\} \mid i = 1, \dots, N, j = 1, \dots, N_i\}$ , such that
$$
h _ {i j} \left(x _ {i j}\right) = \left\{ \begin{array}{l l} \operatorname {R e L U} \left(x _ {i j}\right) & \text {i f} \alpha_ {i j} = 1 \\ \operatorname {I D} \left(x _ {i j}\right) & \text {i f} \alpha_ {i j} = 2 \\ \operatorname {Z e r o} \left(x _ {i j}\right) & \text {i f} \alpha_ {i j} = 3 \end{array} . \right. \tag {3}
$$
The sparsification problem can then be formulated as the following MIP:
$$
\min  _ {\theta = (\omega , \alpha)} \quad \mathcal {L} (\theta) \quad \text {s . t .} \quad \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {N _ {i}} \mathbb {1} \left(\alpha_ {i j} = 1\right) \leq \varepsilon , \tag {4}
$$
where  $\mathbb{1}$  is the indicator function, and the value of  $\varepsilon$  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 4 is hard, as the number of integer variables in  $\alpha$  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  $\pi_{ij}$  indicating the categorical distribution of  $\alpha_{ij}$  assigning to one of the three categories, where  $\pi_{ij}^{k} \triangleq p(\alpha_{ij} = k)$  for  $k = 1,2,3$ . We can then reformulate the problem as:
$$
\min  _ {\omega , \pi} \mathbb {E} [ \mathcal {L} (\theta) ] \quad \text {s . t .} \quad \sum_ {i = 1} ^ {N} \sum_ {j = 1} ^ {N _ {i}} \pi_ {i j} ^ {1} \leq \varepsilon , \quad \alpha_ {i j} \sim \pi_ {i j}, \tag {5}
$$
where  $\pi \triangleq \{\pi_{ij} \mid i = 1, \dots, N, j = 1, \dots, N_i\}$ .
In Equation 5, the sampling procedure  $\alpha_{ij} \sim \pi_{ij}$  is not differentiable. In order to make end-to-end gradient-based optimization possible, we employ the Gumbel-Softmax [28, 46] technique to obtain a continuous approximation of the discrete distribution.
Specifically, for a 3-class categorical distribution  $\pi_{ij}$ , where the class probabilities are denoted as  $\pi_{ij}^{1}, \pi_{ij}^{2}, \pi_{ij}^{3}$ , Gumbel-Max [17] allows us to draw 3-dimensional one-hot categorical samples  $\hat{z}_{ij}$  from the distribution via:
$$
\hat {z} _ {i j} = \operatorname {O n e H o t} \left(\underset {k} {\arg \max } \left(\log \pi_ {i j} ^ {k} + g ^ {k}\right)\right), \tag {6}
$$
where  $g^{k}$  are i.i.d. samples drawn from Gumbel(0,1), which is obtained by sampling  $u^{k} \sim \mathrm{Uniform}(0,1)$  and computing  $g^{k} = -\log (-\log (u^{k}))$ . We can then use the softmax function as a continuous, differentiable approximation of the arg max function:
$$
z _ {i j} ^ {k} = \frac {\exp \left(\left(\log \pi_ {i j} ^ {k} + g ^ {k}\right) / \tau\right)}{\sum_ {k ^ {\prime}} \exp \left(\left(\log \pi_ {i j} ^ {k ^ {\prime}} + g ^ {k ^ {\prime}}\right) / \tau\right)}. \tag {7}
$$
We denote this operation as  $z_{ij} \sim \mathrm{Concrete}(\pi_{ij},\tau)$  [46], where  $\tau$  is a temperature parameter controlling how close the softmax approximation is to the discrete distribution. As the temperature  $\tau$  approaches zero, samples from the Gumbel-Softmax distribution become one-hot and identical to the original categorical distribution.
After obtaining  $z_{ij}$ , we can calculate the activation value  $x_{ij}^{+}$  as a weighted sum of different functional choices:
$$
x _ {i j} ^ {+} = \hat {h} _ {i j} \left(x _ {i j}\right) \triangleq z _ {i j} ^ {1} \cdot \operatorname {R e L U} \left(x _ {i j}\right) + z _ {i j} ^ {2} \cdot \operatorname {I D} \left(x _ {i j}\right) + z _ {i j} ^ {3} \cdot \operatorname {Z e r o} \left(x _ {i j}\right), \tag {8}
$$
and then use gradient descent to optimize both the weight parameters  $\omega$  and the architecture distribution parameters  $\pi$ .
During training, we can also constrain  $z_{ij}$  to be one-hot vectors by using arg max, but use the continuous approximation in the backward pass by approximating  $\nabla_{\theta}\hat{z}_{ij}\approx \nabla_{\theta}z_{ij}$ . This is denoted as "Straight-Through" Gumbel Estimator in [28].
3.4 Optimization algorithm
Instead of limiting the number of regular RLUs from the very beginning of the training process, we start with a randomly initialized neural network and use gradient descent to optimize  $\omega$  and  $\pi$  by minimizing the following objective function until convergence:
$$
\mathbb {E} [ \mathcal {L} (\theta) ] + \lambda R (\pi), \tag {9}
$$
where the regularization term  $R(\pi) \triangleq \sum_{i=1}^{N} \sum_{j=1}^{N_i} \pi_{ij}^1$  aims to explicitly reduce the use of the regular ReLU function. One could also consider adjusting it to  $R(\pi) \triangleq \sum_{i=1}^{N} \sum_{j=1}^{N_i} (\pi_{ij}^1 + \lambda_{\mathrm{ID}} \pi_{ij}^2)$  with a small  $\lambda_{\mathrm{ID}}$  to discourage the use of identity mappings at the same time.
We then take an iterative approach by starting with a relatively large  $\varepsilon_{1}$  and gradually decrease its value for  $K$  iterations with  $\varepsilon_{1} > \varepsilon_{2} > \dots >\varepsilon_{K} = \varepsilon$ . Within each optimization iteration using  $\varepsilon_{k}$ , we first rank the neurons according to  $\max (\pi_{ij}^{2},\pi_{ij}^{3})$  in descending order, and assign the activation function for the top-ranked neurons as ID if  $\pi_{ij}^{2}\geq \pi_{ij}^{3}$ , or Zero otherwise, while keeping the bottom  $\varepsilon_{k}$  neurons intact using Gumbel-Softmax as described in Section 3.3. Subsequently, we continue optimizing  $\omega$  and  $\pi$  using gradient descent to minimize Equation 9. The sparsification process generates a range of models at various sparsification levels for subsequent investigations.
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  _ {u} \quad \sum_ {t} c \left(y _ {t}, u _ {t}\right) \quad \text {s . t .} \quad y _ {t + 1} = \hat {f} _ {\theta} \left(y _ {t ^ {\prime}: t}, u _ {t ^ {\prime}: t}\right), \tag {10}
$$
where  $c$  is the cost function. When the transition function  $\hat{f}_{\theta}$  is a 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-predictive path integral (MPPI) [67] for online planning. Others have also tried applying gradient descent to derive the action signals [36, 37]. However, the number of required samples grows exponentially with the number of inputs and trajectory length. Gradient descent can also be stuck in local optima, and it is also hard to assess the optimality or robustness of the derived action sequence using these methods.
Figure 2: Recover the ground truth piecewise affine functions from data. We evaluate our sparsification pipeline on two hand-designed piecewise affine functions composed of four linear pieces. Our pipeline successfully 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 network is a piecewise affine function, we can formulate Equation 10 as MIP. We assign to each ReLU a binary variable  $a = \mathbb{1}(x\geq 0)$  to indicate whether the pre-activation value is larger or smaller than zero. Given lower and upper bounds on the input  $l\leq x\leq u$  (which we calculate by passing the offline dataset through the sparsified neural networks), the equality  $x^{+} = \operatorname {ReLU}(x)\triangleq \max (0,x)$  can be modeled through the following set of mixed-integer linear constraints:
$$
x ^ {+} \leq x - l (1 - a), \quad x ^ {+} \geq x, \quad x ^ {+} \leq u a, \quad x ^ {+} \geq 0, \quad a \in \{0, 1 \}. \tag {11}
$$
If only a few RLUs are left in the model, Equation 10 can be efficiently solved to global optimality.
The formulation in Equation 11 is the simplest mixed-integer encoding of a ReLU network, and a 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 can lead to tighter convex relaxations of our problem and allow us to identify high-quality solutions of Equation 10 earlier in the branch-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 we typically expect a smaller simulation error from a 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 RLUs are generally less accurate but permit the use of more advanced optimization 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 to counteract prediction errors via online modifications of the action sequence. The iterative 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.
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 specified in different observation spaces, including state, keypoints, and object-centric representations.