Data Center Chiller Plant Optimization via Mixed-Integer Nonlinear Differentiable Predictive Control

The rapid expansion of artificial intelligence and digital infrastructure has led to a sharp increase in global data center electricity demand. In the United States alone, data centers currently consume about 150 TWh of electricity annually and are projected to grow to between 200 TWh and 500 TWh by 2030 (Aljbour et al., 2024). With this expansion comes an accompanying demand for high-efficiency cooling systems in data centers. Chiller plants are commonly deployed to perform this task. A chiller plant typically includes heat pumps, pumps, and cooling towers, with the heat pumps being the most energy-intensive components within the plant. Control of chiller plants involves both continuous decision variables (e.g., chilled water supply temperature, differential pressure) and discrete decision variables (e.g., number of chillers active, pump mode). Optimizing these coupled decisions can significantly reduce electricity consumption, ensure reliable operation, and extend the lifespan of costly mechanical assets. However, this requires solving a high-dimensional mixed-integer nonlinear optimal control problem, which is computationally challenging to solve in real time.

To address these challenges, the proposed approach utilizes Differentiable Predictive Control (DPC) (Drgona et al., 2022), which formulates parametric MPC problems as differentiable programs, enabling offline self-supervised learning of explicit neural control policies through gradient-based learning. Expanding upon this foundation, recent work on Mixed-Integer DPC (MI-DPC) (Boldocky et al., 2025) extended this idea to handle mixed-integer optimal control problems (MI-OCPs) with discrete decision variables.

The continuous-time thermal dynamics of the multi-chiller system are \begin{align} Q^{(i)}(t) &= \eta_\mathrm{r}\dot{m}^{(i)}(t)\; c_\mathrm{p}\delta^{(i)}(t) \left( T_\mathrm{r}(t)-T^{(i)}_\mathrm{s}(t) \right), \label{eq:cooling_delivered}\\ \dot{T}_\mathrm{r}(t) &= \frac{1}{C_\mathrm{r}}\left(\tilde{Q}_\text{load}(t) - \sum_{i=1}^{M}Q^{(i)}(t) \right), \label{eq:temp_return}\\ \dot{T}_\mathrm{s}^{(i)}(t) &= -\frac{c_\mathrm{p} \eta_\mathrm{s}}{C^{(i)}}\left(\delta^{(i)}(t)\dot{m}^{(i)}(t) \left(T^{(i)}_\mathrm{s}(t)-T^{(i)}_\mathrm{e}(t)\right)\right),\!\! \label{eq:temp_supply} \end{align}
where \( Q^{(i)}(t)\!\in\!\mathbb{R}_{\ge0} \) is the delivered cooling from chiller \( i \) at time \( t \), \( \delta^{(i)}(t)\!\in\!\{0,1\} \) denotes the chiller's on/off state, \( c_\mathrm{p}\! \in\!\mathbb{R}_{>0} \) is the specific heat capacity of the coolant, \( T_\mathrm{r}(t)\!\in\!\mathbb{R} \) is the return water temperature, and \( C_\mathrm{r},C^{(i)}\!\in\!\mathbb{R}_{>0} \) are the thermal capacitances of the return and chiller loops, \( T_\mathrm{s}^{(i)}(t)\!\in\!\mathbb{R} \) is the supply water temperature for chiller \( i \in \{1, ..., M\} \). Moreover, \( \eta_\mathrm{r},\eta_\mathrm{s}\!\in\!(0,1) \) denote the heat-exchanger efficiency coefficients, and \( \tilde{Q}_\text{load}(t)\!\in\! \mathbb{R}_{\ge0} \) is the effective thermal load, estimated using a low-pass, discrete-time finite impulse response (FIR) filter \( \tilde{Q}_\text{load} (t) = \sum_{l=0}^{L} h_l Q_\text{load}(t-l\Delta t) \), where \( Q_\text{load}(t)\in\mathbb{R}_{\ge 0} \) is the direct server load; an exogenous variable in this context. In the following, the continuous-time dynamics are discretized using a fourth-order Runge-Kutta method with a time-step length \( \Delta t \) and discrete time index \( k \); denoted as \( f(x_k, u_k, \tilde{Q}_{\text{load},k}, \Delta t) : \mathbb{R}^{n_x} \times \mathbb{R}^{n_u} \times \mathbb{R} \rightarrow \mathbb{R}^{n_x} \), where \( x_k = [T_r(k), T_s^{(1)}(k), ..., T_s^{(M)}(k)]^T \).

The power consumption of chillers \( P^{(i)}_\text{chiller}(t)\in\mathbb{R}_{\ge 0} \) and pumps \( P^{(i)}_\text{pump}(t)\in\mathbb{R}_{\ge 0} \) are modeled as \begin{align} P^{(i)}_\text{chiller}(t) &= \frac{Q^{(i)}(t)}{\text{COP}^{(i)}(t)}+\rho^{(i)}\delta^{(i)}(t) \label{eq:chiller_power},\\ \text{COP}^{(i)}(t) &= a_0^{(i)}+a_1^{(i)}\frac{Q^{(i)}(t)}{Q_\text{max}^{(i)}}+a_2^{(i)}\left(\frac{Q^{(i)}(t)}{Q_\text{max}^{(i)}}\right)^2,\\ P^{(i)}_\text{pump}(t) &= \gamma^{(i)}\left( \dot{m}^{(i)}(t)\delta^{(i)}(t) \right)^3, \label{eq:pump_power} \end{align} for the system constants \( a_{0:2}^{(i)},\gamma^{(i)},\rho^{(i)},Q_\text{max}^{(i)}\!\in\!\mathbb{R} \). The constant \( \rho \) represents the base power required for chiller operation. In this formulation, the chiller's coefficient of performance \( \text{COP}^{(i)}(t)\in\mathbb{R}_{\ge0} \) is a function of the part-load ratio (PLR), the ratio of current cooling to maximum cooling capacity. Ensuring that each chiller is operating near its highest efficiency load ratio is the core of the underlying optimal control problem.

We formulate the mixed-integer optimal control problem as a differentiable program:
where, the parameters \( \theta_{(\cdot)} \) of the control policies \( \pi_{\theta_1} \), \( \pi_{\theta_2} \), and \( \pi_{\theta_3} \) are the optimization arguments. The policies map the vector of control parameters \( \xi \) to continuous-valued control variables \( \dot m \) (mass flow), \( T_\mathrm{e} \) (evaporation temperature), and a slack variable \( \tilde\delta\!\in\!\mathbb{R}^{M-1} \) (relaxed value of the binary control variable \( \delta \)). By fixing one of the binary variables, the dimension of \( \tilde\delta \) is reduced to \( M\!-\!1 \). This is done as a safety mechanism to ensure that at least one chiller is always on. To compute control variables, separate neural modules are used, which are represented by fully connected deep neural networks. Moreover, the control parameters are sampled from a known probability distribution \( P_\xi \) to emulate scenarios that are likely to occur during the system's operation. \(N\) denotes prediction horizon length, whereas functions \( q(\cdot) \) and \( p(\cdot) \) denote state and input penalty functions weighted by \( \lambda_x \) and \( \lambda_u \), respectively.

The binary integrality is enforced using a threshold-based discretization scheme with a fixed threshold of \(0.5\), given as \begin{align} \delta^{(i)}_{k} = & \begin{cases} 1, & \text{if } \tilde\delta^{(i)}_{k}>0.5,\\[0pt] 0, & \text{otherwise.}\\[0pt] \end{cases} \end{align} This function behaves analogously to a Heaviside function exhibiting a discontinuity at \(0.5\). To enable gradient-based optimization, a Straight-through estimator heuristic by (Bengio et al., 2013), is used. The gradient is approximated using its differentiable surrogate, given by a scaled sigmoid function \( \sigma(x)\!=\!\left(1+e^{-\mu (x-0.5)} \right)^{-1} \), where \( \mu\!\in\!\mathbb{R}_{>0} \) is a slope parameter determining the sharpness of the transition; additionally, the function is centered at \(0.5\). Accordingly, the derivative is approximated as \begin{align} \frac{\partial \delta^{(i)}_{k}}{\partial \tilde{\delta}^{(i)}_{k}} \approx \frac{\partial \sigma\!\left(\tilde{\delta}^{(i)}_{k}\right)}{\partial \tilde{\delta}^{(i)}_{k}}. \end{align}
The differentiable computational graph, comprises four components: a continuous neural policy, a differentiable rounding layer, a differentiable ordinary differential equation (ODE) solver, and a model predictive loss function. Given synthetically sampled control parameters that include the initial conditions and a length-\(N\) vector of exogenous variables (server load), a closed-loop rollout of length \(N\) is performed through the integrated system dynamics, constituting the single-shooting forward pass. Consequently, the loss function is evaluated based on the resulting closed-loop trajectories. The backward pass is used to enable policy optimization by propagating gradients through the computational graph. Specifically, gradients of the loss function with respect to policy parameters are obtained by back-propagating through the integrated system dynamics using either the backpropagation through time algorithm (Puskorius and Feldkamp, 1994).

To promote binary polarization of the relaxed variable \( \tilde \delta \) and to complement the \( \Delta\delta \) penalty in suppressing high-frequency switching, a binary-variance regularization (BVR) is introduced as \( \|\tilde\delta_k\odot(1-\tilde\delta_k)\|_\Lambda^2 \), where \( \Lambda\!\in\!\mathbb{R}^{M-1\times M-1} \) denotes the weighting matrix. As the weight increases, the regularization penalizes values of \( \tilde\delta \) near the discretization threshold (\(0.5\)), as well as those outside the \([0,1]\) range. Based on extensive numerical experiments, this regularization has proven crucial to the proposed framework, as its omission leads to \( \tilde\delta \) values clustering around the discretization threshold, thereby increasing the likelihood of high-frequency switching behavior. The influence of this regularization on the relaxed binary variable for different penalty magnitudes is illustrated in the following Figure.