# Moment Constrained Optimal Transport for Energy Demand Management of Heterogeneous Loads

Julien Cardinal, Thomas Le Corre, Ana Bušić

► To cite this version:

Julien Cardinal, Thomas Le Corre, Ana Bušić. Moment Constrained Optimal Transport for Energy Demand Management of Heterogeneous Loads. NETGCOOP 2025 - 12th International Conference of Networks, Games, Control and Optimization, Oct 2025, Bilbao, Spain. hal-05375478

HAL Id: hal-05375478

https://hal.science/hal-05375478v1

Submitted on 27 Nov 2025

HAL is a multi-disciplinary open access
archive for the deposit and dissemination of sci-
entific research documents, whether they are pub-
lished or not. The documents may come from
teaching and research institutions in France or
abroad, or from public or private research centers.

L'archive ouverte pluridisciplinaire HAL, est
destinée au dépôt et à la diffusion de documents
scientifiques de niveau recherche, publiés ou non,
émanant des établissements d'enseignement et de
recherche français ou étrangers, des laboratoires
publics ou privés.

<figure><img src="image_1.png" /></figure>

Distributed under a Creative Commons Attribution 4.0 International License
# Moment Constrained Optimal Transport for Energy Demand Management of Heterogeneous Loads

Julien Cardinal¹,²[0009-0006-7431-2762], Thomas Le Corre¹,²[0000-0003-1471-6053], and Ana Bušić¹,²[0000-0002-4133-3739]

¹ Inria, Paris, France

² DI ENS, ENS, PSL University, Paris, France

**Abstract.** This paper addresses the problem of coordinating a large population of heterogeneous electrical loads, such as electric vehicles (EVs) and water heaters (WHs), under global operational constraints. We extend the Moment Constrained Optimal Transport for Control (MCOT-C) framework to accommodate multiple classes of agents with distinct dynamics and cost structures. Our formulation relies on a mean-field limit that captures agent heterogeneity through class-specific distributions. We propose a scalable gradient descent algorithm and a Model Predictive Control (MPC) scheme that enables online adaptation of this algorithm to uncertain or progressively revealed agent information. The proposed approach is validated through numerical experiments on real datasets [8,9] for EVs and WHs, demonstrating the effectiveness of this method in enforcing global constraints while preserving agent-level dynamics.<fnref n="3" />

**Keywords:** Mean Field Control · Optimal Transport · Smart Grids.

## 1 Introduction

Demand-side management (DSM) has emerged as a key strategy for enhancing the flexibility and reliability of modern power systems, particularly in the context of increasing penetration of renewable energy sources. Unlike traditional generation-side methods, DSM enables system operators to influence and coordinate the consumption patterns of end-users in real-time or near-real time. This is achieved by leveraging the flexibility of a wide range of distributed energy resources (DERs). The aggregated control of such devices can provide valuable services to the grid, including load shifting, peak shaving, frequency regulation, and congestion management, thus contributing to system stability and economic

<footnote>This work was carried out in the framework of the AI-NRGY project, funded by France 2030 (Grant No: ANR-22-PETA-0004).</footnote>

<footnote n="3">³ The code used to generate the results of this paper is available at https://github.com/Kreyparion/Heterogeneous-MCOT.</footnote>
efficiency [5][6]. In demand-side management applications, the number of con-
trollable devices is typically very large, which naturally motivates the use of
mean-field control (MFC) approaches [1].

MFC problems are well-studied because of their versatility in domains rang-
ing from economics and energy management to large-scale networked infrastruc-
tures [1,2]. When a large population of agents interacts through a common cost,
the difficulty of solving the associated control problem increases with the number
of agents. The mean-field approximation mitigates this difficulty by represent-
ing the system through a single, representative agent influenced by the aggregate
state of the population, leading to decentralized control strategies that align in-
dividual behavior with a shared global objective.

In demand-side management control problems, the populations of devices
are highly heterogeneous, both in their physical nature (such as electric ve-
hicles, refrigerators, or water heaters), and in their individual characteristics
(for instance differing battery capacities or charging rates in the case of electric
vehicles). Blindly applying homogeneous decentralized strategies may result in
suboptimal or even infeasible behavior [3]. To efficiently control heterogeneous
agents, structured population models [1] or heterogeneous mean-field games [3]
have been proposed. The challenge of controlling distributed energy resources
has been addressed in [4], where a duality-based approach is employed, closely
related to the one considered in this work.

Optimal transport (OT) provides an intuitive approach for modeling the evo-
lution of the aggregate distributions considered in MFC. Moment-Constrained
Optimal Transport (MCOT) has been proposed as a relaxation of the classical
OT problem, where the marginal constraints are relaxed using moment inequal-
ities [7]. This relaxation was first introduced to approximate the optimal trans-
port problem but was recently considered as a formulation of a mean field control
problem, for a homogeneous population [11]. In this work, we extend the MCOT
framework to incorporate agent heterogeneity, by formulating a heterogeneous
moment-constrained optimal transport problem.

**Contributions:** We develop a tractable optimization framework for the real-time coordination of heterogeneous loads, combining ideas from optimal transport and mean field control. The first contribution of this paper is to extend the framework of Moment Constrained Optimal Transport for Control (MCOT-C) [11] to a heterogeneous setting. The second contribution is to develop a Model Predictive Control (MPC) approach where the agents’ data is progressively discovered during the day. The relevance of this method is checked on a real data set of Water Heaters and Electric Vehicles.

**Organisation of the paper:** Section 2 introduces the formulation of the problem of controlling heterogeneous agents as a Moment Constrained Optimal Transport problem. The main results on the existence of solutions and expression of the dual and its gradients are given in this section, leading to a gradient descent algorithm. In Section 3, we propose an MPC version in which the data is discovered sequentially. An application illustrated this method for two types of agents: Electric Vehicles (dataset [8]) and Water Heaters (dataset [9]).
# 2 Heterogeneous MCOT

## 2.1 Heterogenous Control Problem

We consider a population of $N$ agents, denoted $\{X^i\}_{1\le i\le N}$, partitioned into $H$ classes. Each class $h \in \{1, \dots, H\}$ includes a subset $\mathcal{I}^{(h)} \subseteq \{1, \dots, N\}$ of agents sharing common dynamics and cost. Let $N^{(h)} = |\mathcal{I}^{(h)}|$ denote the number of agents in class $h$, with $\sum_{h=1}^H N^{(h)} = N$. The proportion of each class is denoted $\alpha^{(h)}: \forall h \in \{1, \dots, H\}, \alpha^{(h)} = \frac{N^{(h)}}{N}$.

A central agent coordinates the agents to satisfy a set of $A$ global constraints:

$$ \forall a \in \{1, \dots, A\}, \quad \sum_{h=1}^{H} \sum_{i \in \mathcal{I}^{(h)}} f_{a}^{(h)}(X^i) \le 0, \quad (1) $$

where $f^{(h)}: \mathcal{X}^{(h)} \to \mathcal{R}^A$ is a function depending on the class $h$.

An example of such a constraint is a total power consumption cap $P_{\max}$ that must not be exceeded at any time. Equality constraints, such as signal tracking, can be formulated as pairs of inequality constraints and are treated similarly, as in Section 3. Each agent is defined by their state $\{s^i\} \in S$ and makes a control decision $\{w^i\} \in W$, which induces a cost $c^{(h)}$ (examples of this cost are provided in Section 3). The central agent aims at minimizing the following problem:

$$ \min_{\{w^i\}_{1\le i\le N}} \left\{ \sum_{h=1}^{H} \sum_{i \in \mathcal{I}^{(h)}} c^{(h)}(s^i, w^i) : \sum_{i=1}^{N} f(s^i, w^i) \le 0 \right\} \quad (2) $$

This formulation sets the stage for a class-wise mean field approach to control design, developed in the following sections.

## 2.2 Heterogeneous MCOT

We introduce in this section the Mean Field formulation [10] of Problem 2 to the case of $H$ distinct classes of agents. Each class $h \in \{1, \dots, H\}$ is associated with a population of agents sharing common nominal dynamics and costs. The behavior of class-$h$ agents is described by a common distribution $\mu_1^{(h)} \in \mathcal{B}(\mathcal{X}^{(h)})$, where $\mathcal{B}(\mathcal{X}^{(h)})$ denotes the Borel probability measures on $\mathcal{X}^{(h)} = S^{(h)} \times W^{(h)}$, the product space of the state and control spaces. The aggregate system is thus described by the collection $\{\mu_1^{(h)}\}_{h=1}^H$, each representing a homogeneous subpopulation. The goal is to control each class by modifying its distribution, resulting in a transport of mass between distributions over $\mathcal{X}^{(h)}$.

We aim to identify joint distributions $\pi^{(h)} \in \mathcal{B}(\mathcal{X}^{(h)} \times \mathcal{X}^{(h)})$ for each class $h$, representing the coupling between initial and controlled trajectories. Each $\pi^{(h)}$ is subject to constraints reflecting both the nominal dynamics and control objectives specific to class $h$ and is associated with a cost: $\langle \pi^{(h)}, c^{(h)} \rangle = \int c^{(h)}(x, y)\pi^{(h)}(x, y)dxdy$. We define the marginals $\pi_1^{(h)} = \int \pi^{(h)}(x, y)dy$ and $\pi_2^{(h)} = \int \pi^{(h)}(x, y)dx$, with $\pi_1^{(h)} = \mu_1^{(h)}$ fixed. Second marginals $\pi_2^{(h)}$ must belong to the following set of distributions satisfying the global moment constraints:

$$ \mathcal{P}_f = \left\{ \{\mu^{(h)}\}_{h=1}^H \in \prod_{h=1}^H \mathcal{B}(\mathcal{X}^{(h)}) : \forall a \in \{1, \dots, A\}, \sum_{h=1}^H \alpha^{(h)} \langle \mu^{(h)}, f_a^{(h)} \rangle \le 0 \right\} \quad (3) $$
We define the feasible coupling set, as bivariate distributions having for first marginal $\mu_1^{(h)}$ and whose two marginals coincide on the state space $S$:

$$ \mathcal{U}(\mu_1^{(h)}) = \left\{ \begin{array}{l} \pi^{(h)} \in \mathcal{B}(\mathcal{X}^{(h)} \times \mathcal{X}^{(h)}) : \pi_1^{(h)} = \mu_1^{(h)} \\ \text{and } \forall s_1 \neq s_2 \in S, \pi^{(h)}((s_1, .), (s_2, .)) = 0 \end{array} \right\} \quad (4) $$

An entropic regularizer is introduced for computational reasons and to impose hard constraints on agent behavior [11]:

$$ D_{\mathrm{KL}}(\pi^{(h)} \| \mu_1^{(h)} \otimes \mu_2^{(h)}) = \int \log \left( \frac{\pi^{(h)}(x, y)}{\mu_1^{(h)}(x) \mu_2^{(h)}(y)} \right) \pi^{(h)}(x, y) \, dx dy \quad (5) $$

leading to the formulation of the problem: **Multi-Class MCOT-C**

$$ \begin{array}{l} \min_{\{\pi^{(h)}\}_{h=1}^H} \sum_{h=1}^H \alpha^{(h)} \left( \langle \pi^{(h)}, c^{(h)} \rangle + \varepsilon D_{\mathrm{KL}}(\pi^{(h)} \| \mu_1^{(h)} \otimes \mu_2^{(h)}) \right) \quad (6) \\ \text{s.t. } \forall h, \pi^{(h)} \in \mathcal{U}(\mu_1^{(h)}), \{\pi_2^{(h)}\}_{h=1}^H \in \mathcal{P}_f \end{array} $$

## 2.3 Dual Problem

In this subsection, we begin by stating the assumptions, then present the dual formulation of the heterogeneous MCOT problem 6. We then show that this dual problem admits a closed-form expression (Proposition 1), which is differentiable, and we provide the explicit formula for its gradient (Proposition 2).

**Assumptions:** In the following, we make the following assumptions,

(A1) Each $c^{(h)}: \mathcal{X}^{(h)} \times \mathcal{X}^{(h)} \to \mathbb{R}_+$ and $f^{(h)}: \mathcal{X}^{(h)} \to \mathbb{R}^{A^{(h)}}$ are continuous. For each class $h$, $\mathcal{P}_{f,r}^{(h)} = \{\mu \in \mathcal{B}(\mathcal{X}^{(h)}) : \forall a \in \{1, \dots, A\}, \sum_{h=1}^H \langle \mu, f_a^{(h)} \rangle \le r\}$ is non-empty for all $r$ in a neighborhood of 0.

(A2) Each $\mu_1^{(h)}, \mu_2^{(h)}$ has compact support. The feasibility property under perturbations holds for each class.

(A3) The covariance $\Sigma^{(h)}$ of $Y^{(h)}$ is positive definite for $Y^{(h)} \sim \mu_2^{(h)}$.

**Dual Problem:** We denote the Lagrange multiplier associated to the constraints 1 by $\lambda$ and we derive the dual problem from 6: $\forall \lambda \in \mathbb{R}^A$,

$$ \begin{array}{l} \varphi(\lambda) = \min_{\{\pi^{(h)}\}_{h=1}^H} \sum_{h=1}^H \alpha^{(h)} \left( \varepsilon D_{\mathrm{KL}}(\pi^{(h)} \| \mu_1^{(h)} \otimes \mu_2^{(h)}) - \langle \pi^{(h)}, \ell_{\lambda}^{(h)} \rangle \right) \quad (7) \\ \text{s.t. } \forall h, \pi^{(h)} \in \mathcal{U}(\mu_1^{(h)}) \end{array} $$

with $\ell_{\lambda}^{(h)}(x, y) = -\lambda^{\top} f^{(h)}(y) - c^{(h)}(x, y)$. For each class $h \in \{1, \dots, H\}$, each $\lambda \in \mathbb{R}_+^A$, $\varepsilon > 0$, and $x = (s, w_y) \in S^{(h)} \times W^{(h)} = \mathcal{X}^{(h)}$, we define:

$$ B_{\lambda, \varepsilon}^{(h)}(x) = \varepsilon \log \int \exp \left( \varepsilon^{-1} \ell_{\lambda}^{(h)} ((s, w_x), (s, w_y)) \right) \mu_2^{(h)}(s, w_y) dw_y \quad (8) $$
**Proposition 1.** (i) The minimum of the dual problem gives:

$$ \forall \lambda \in \mathbb{R}^A, \varphi(\lambda) = - \sum_{h=1}^{H} \alpha^{(h)} \langle \mu_1^{(h)}, B_{\lambda, \varepsilon}^{(h)} \rangle. \quad (9a) $$

(ii) The optimal coupling $\pi^{(h)}$ satisfies: $\forall x = (s_x, w_x), y = (s_y, w_y) \in \mathcal{X}^{(h)}$

$$ \pi_{\lambda}^{(h)}(x, y) = \mu_{1}^{(h)}(x)\mu_{2}^{(h)}(y)\delta_{s_x}(s_y) \exp\left(L_{\lambda}^{(h)}(x, y)\right), \quad (9b) $$

where: $L_{\lambda}^{(h)}(x, y) = \varepsilon^{-1} \{\ell_{\lambda}^{(h)}(x, y) - B_{\lambda, \varepsilon}^{(h)}(x)\}$

(iii) There is no duality gap.

*Proof.* As each term in the sum in (7) only depends of $\pi^{(h)}$, the minimization can be performed term by term:

$$ \varphi(\lambda) = \sum_{h=1}^{H} \alpha^{(h)} \min_{\pi^{(h)}} \left( \varepsilon D_{\text{KL}}(\pi^{(h)} \| \mu_{1}^{(h)} \otimes \mu_{2}^{(h)}) - \langle \pi^{(h)}, \ell_{\lambda}^{(h)} \rangle \right) $$

The minimization subproblems were solved in [11] which gives (i),(ii) and (iii). The proof is based on convex duality between relative entropy and log moment generating functions. □

We now introduce the scaled dual function and study its derivative: $\mathcal{J}(\zeta) := -\varepsilon^{-1}\varphi(\varepsilon\zeta)$.

**Proposition 2.** Function $\mathcal{J}^{(h)}$ is convex and continuously differentiable, with:

$$ \nabla \mathcal{J}(\zeta) = \sum_{h=1}^{H} \alpha^{(h)} \langle \mu_{\lambda}^{(h)}, f^{(h)} \rangle = \sum_{h=1}^{H} \alpha^{(h)} \mathbb{E}_{\lambda}^{(h)}[f^{(h)}] \quad (10a) $$

$$ \nabla^2 \mathcal{J}(\zeta) = \sum_{h=1}^{H} \alpha^{(h)} \mathbb{E}_{\lambda}^{(h)} \left[ f^{(h)}(Y) f^{(h)}(Y)^{\top} - \mathbb{E}_{\lambda}^{(h)} [f^{(h)}(Y) | X] \mathbb{E}_{\lambda}^{(h)} [f^{(h)}(Y) | X]^{\top} \right] \quad (10b) $$

where $\mu_{\lambda}^{(h)}$ is the second marginal of $\pi_{\lambda}^{(h)}$ and $\mathbb{E}_{\lambda}^{(h)}$ is the expectancy according to $\mu_{\lambda}^{(h)}$. It implies that $\mathcal{J}$ is strictly convex.

*Proof.* The proof is a direct differentiation of the dual expressions in Proposition 1. □

## 2.4 Algorithm

We design a gradient descent algorithm (Algorithm 1), in which $\mathbb{E}_{\lambda}^{(h)}[f^{(h)}]$ is computed for each class either directly when the state space associated with this class is finite and not too large (as in Section 3.2) or through Monte

**Algorithm 1** a Gradient Descent algorithm to solve MCOT

1: Initialize $\lambda$

2: **for** each $k = 1, \dots, K$ **do**

3: $G \leftarrow \sum_{h=1}^{H} \alpha^{(h)} \mathbb{E}_{\lambda}^{(h)}[f^{(h)}]$

4: $\lambda \leftarrow \lambda - \rho_k G$

5: **end for**

Carlo Methods when this state space is too large or even infinite (as in Section 3.3). $\{\rho_k\}_{k=1}^K$ the step size and K the number of gradient descent iterations are chosen empirically to accelerate convergence.
# 3 Model Predictive Control

## 3.1 The Model Predictive Control Algorithm

**Algorithm 2** Model Predictive Control for Heterogeneous loads

1: Initialize $G^{MPC}, \lambda \in \mathbb{R}^T$
2: **for** each $t = 1, \dots, T$ **do**
3: &nbsp;&nbsp;&nbsp;&nbsp;**for** each $k = 1, \dots, K$ **do**
4: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$$G \leftarrow \sum_{h=1}^{H} \alpha^{(h)} \mathbb{E}_{\lambda_{t:T}}^{(h)}[f_{t:T}^{(h)}]$$
5: &nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;&nbsp;$$\lambda_{t:T} \leftarrow \lambda_{t:T} - \rho_k G$$
6: &nbsp;&nbsp;&nbsp;&nbsp;**end for**
7: $$G_t^{MPC} \leftarrow \sum_{h=1}^{H} \alpha^{(h)} \mathbb{E}_{\lambda_{t:t}}^{\text{test},(h)}[f_{t:t}^{(h)}]$$
8: **end for**

In the Model Predictive Control (MPC) setting (Algorithm 2), we aim at satisfying some predefined constraints (e.g. tracking or maximum consumption constraints) using prediction of the future dynamic of the loads. For the Water Heaters, the unknowns are the future water drains and for the Electric Vehicles, it is the number of EVs and their charging needs. These unknowns are progressively discovered by the algorithm at each time $t$, which updates the density considered in the expectancy $\mathbb{E}_{\lambda}^{(h)}$. More details on this update of the density for each type of agent are presented in the next sections. We denote $f_{t_1:t_2}^{(h)}$, $\lambda_{t_1:t_2}$ the global constraints, respectively the slice of $\lambda$ from $t_1$ to $t_2$, and $T$ the number of time intervals.

## 3.2 Electric Vehicles

As in [11], we consider a population of Electric Vehicles with their behavior derived from the Elaad OpenDataset [8]. This dataset is composed of 10.000 random transactions from public charging stations operated by EVnetNL in the Netherlands, in the year 2019. We limit ourselves to a typical workday, discretized in $T = 144$ intervals of size $\delta t = 10$ minutes, and separate the dataset between a training set $D_{train}$ and a test set $D_{test}$.

Every vehicle is described by its arriving time $t_a$, its leaving time $t_l$ (which is also the time limit for which it must be fully charged) and its charging needs $\Delta t_n$, the duration it needs to be fully charged. Therefore, the state space of all EVs is written as follows:

$$S^{\text{EVs}} = \underbrace{\{1, \dots, T\}}_{\text{Arriving time } t_a} \times \underbrace{\{1, \dots, T\}}_{\text{Leaving time } t_l} \times \underbrace{\{1, \dots, T\}}_{\text{Charging need } \Delta t_n} \quad (11)$$

We consider here that all EVs have the same power rating, but with different battery levels, represented by the charging needs $\Delta t_n$.

Furthermore, we consider that we can control the charging starting time $t_c$ of every EV. At a given time step $t \in \{1, \dots, T\}$, we can define the charging starting time, anywhere in $W_t^{\text{EVs}} = \{t, \dots, T\}$ if it also satisfies the constraints: $t_a \le t_c \le t_l - \Delta t_n$. Finally, we define $\mathcal{X}_t^{\text{EVs}} = S^{\text{EVs}} \times W_t^{\text{EVs}}$ the controlled state space. We can now define distributions over $S^{\text{EVs}}$ and $\mathcal{X}_t^{\text{EVs}}$:

- We define $\nu$ as the fraction of vehicles in each state in the training set.

$$\forall s \in S^{\text{EVs}}, \quad \nu(s) = \frac{\#\{s' \in D_{\text{train}} : s' = s\}}{\#D_{\text{train}}} \quad (12)$$
We denote as {$s_t^i$} $\subset D_{test}$ the available vehicles that are present and not charging yet.

- We define as $\nu_t$, the empirical fraction of vehicles left to charge:

$$ \forall s = (t_a, t_l, \Delta t_n) \in S^{\text{EVs}}, \quad \nu_t(s) = \begin{cases} \frac{1}{N_t} \sum_i \delta(s - s_t^i) & \text{if } t_a \le t \\ \frac{N^{\text{EVs}}}{N_t} \nu(s) & \text{if } t_a > t \end{cases} \quad (13) $$

With $N_t = \int_{S^{\text{EVs}}} \sum_i \delta(s - s_t^i) ds + N^{\text{EVs}} \int_{S^{\text{EVs}}} \nu(s) \mathbf{1}_{t_a > t}(s) ds$ the normalization factor, counting the number of vehicles already arrived and not charging plus the number of vehicles that are estimated to arrive.

- The nominal distribution $\mu_{1,t}$ is defined by the "Plug when Arrive" strategy:

$$ \forall (s, t_c) = (t_a, t_l, \Delta t_n, t_c) \in \mathcal{X}_t^{\text{EVs}}, \quad \mu_{1,t}(s, t_c) = \nu_t(s) \delta(t_c - t_a) \quad (14) $$

- $\mu_{2,t}$ is defined as the "Plug with a uniform distribution" strategy: $\forall s = (t_a, t_l, \Delta t_n, t_c) \in \mathcal{X}_t^{\text{EVs}}$,

$$ \mu_{2,t}(s, t_c) = \begin{cases} \frac{1}{t_l - \Delta t_n - \max(t_a, t) + 1} \nu_t(s) & \text{if } \max(t_a, t) \le t_c \le t_l - \Delta t_n \\ 0 & \text{otherwise} \end{cases} \quad (15) $$

For the MPC algorithm, we define:

- The local cost function $c$ as a quadratic penalization:

$$ c: \mathcal{X}_t^{\text{EVs}} \times S^{\text{EVs}} \times W_t^{\text{EVs}} \to \mathbb{R} \\ x, (t_a, t_l, \Delta t_n), t_c \mapsto (t_c - t_a)^2 \quad (16) $$

$c$ is a penalty for starting charging long after the vehicle arrives.

- We consider that all EVs have the same power consumption of 6kW and define the power (in kW) of one as: $P^{\text{EVs}}: \mathcal{X}_t^{\text{EVs}} \to \mathbb{R}^T$, such that $\forall t' \in \{t, \dots, T\}$:

$$ P_{t'}^{\text{EVs}}: \begin{cases} \mathcal{X}_t^{\text{EVs}} \\ t_a, t_l, \Delta t_n, t_c \end{cases} \to \begin{cases} \mathbb{R} \\ 6 \text{ if } t_c \le t' < t_c + \Delta t_n \\ 0 \text{ otherwise} \end{cases} \quad (17) $$

## 3.3 Water Heaters

As in [12], we consider a large population of Water Heaters with their water drains drawn in a dataset from the SMACH (Multi-agent Simulation of Human Activity in the Household) platform [9]. We limit the dataset to working days and separate the dataset between a training set $D_{train}$ and a test set $D_{test}$.

At time $t \in [0, 24]$, the i-th WH is modeled by its mean temperature $\theta^i(t) \in \Theta$ and its power mode $m^i(t) \in \{0, 1\}$ (Off/On). These WHs must then follow the Ordinary Differential Equation (ODE):

$$ \frac{d\theta(t)}{dt} = -\underbrace{\rho(\theta(t) - \theta_{amb})}_{\text{heat loss}} + \underbrace{\sigma m(t)p}_{\text{Joule effect}} - \underbrace{\sigma\epsilon(t)}_{\text{water drain}} \quad (18) $$
with $\rho$ the fraction of heat loss by minute, $\sigma$ the specific heat capacity of the volume of water, $p$ the heating power, $\theta_{amb}$ the room temperature, and $\epsilon(t)$ the power equivalent of the water drains at time $t$.

We are interested in controlling their behavior over a day, discretized in $T = 144$ intervals of size $\delta t = 10$ minutes. The previous ODE (18) is thus discretized with the following update:

$$ \theta_{t+1} = \theta_t - (\rho(\theta_t - \theta_{amb}) + \sigma m_t p_{max} - \sigma \epsilon_t)\delta t \quad (19) $$

where for $t \in \{0, \dots, T\}$, $\theta_t$, $m_t$ and $\epsilon_t$ are the values of $\theta$, $m$ and $\epsilon$ at time $t \times \delta t$.

We define the nominal dynamics as the uncontrolled dynamic [13,14]. A water heater aims at keeping its mean temperature between $\theta_{min} = 50^{\circ}C$ and $\theta_{max} = 65^{\circ}C$ by turning the water heater Off whenever the temperature reaches $\theta_{max}$ and turning it back On whenever the temperature reaches below $\theta_{min}$. A nominal trajectory with the drains drawn in $\mathcal{D}_{train}$ is represented in Fig. 1(a).

The full update equation for the nominal policy $\mu_1$ can then be written as follows:

$$ \left\{ \begin{array}{l} \theta_{t+1} = \theta_t - \rho \delta t (\theta_t - \theta_{amb}) + \sigma \delta t m_t p_{max} - \sigma \epsilon_t \\ m_{t+1} = \begin{cases} m_t & \text{if } \theta_{t+1} \in [\theta_{min}, \theta_{max}] \\ 0 & \text{if } \theta_{t+1} \ge \theta_{max} \\ 1 & \text{if } \theta_{t+1} \le \theta_{min} \end{cases} \\ \theta_0, m_0 \sim \nu_0 \end{array} \right. \quad (20) $$

with $\nu_0$ being the initial density: the initial temperature $\theta_0$ is uniformly distributed in $[\theta_{min}, \theta_{max}]$ and the initial mode $m_0$ is 1 (On) with probability 0.25 and 0 (Off) with probability 0.75

Therefore, we can define the state space of all WHs as follows:

$$ S^{\text{WHs}} = \underbrace{\Theta}_{\text{Initial Temperature } \theta_0} \times \underbrace{\{0, 1\}}_{\text{Initial mode } m_0} \times \underbrace{\mathbb{R}_+^T}_{\text{Random Drains } \{\epsilon_t\}} \quad (21) $$

Furthermore, we allow the WHs to switch their power mode (from On to Off or from Off to On), while the temperature is still between the two bounds

<figure><img src="image_2.png" /><figcaption>Fig. 1: Example of trajectory for a single water heater, with $\theta_0 = 54^{\circ}C$ and $m_0 = 0$</figcaption></figure>
$\theta_{t+1} \in [\theta_{min}, \theta_{max}]$. We limit ourselves to two additional switches per day per WH, compared to its nominal dynamics. We will note these two times $t_1, t_2 \in \{1, \dots, T\}^2$. This limitation avoids frequent switching, which may be undesirable. Fig. 1(b) shows an example of a trajectory with two switches.

Therefore, the power mode update for the policy $\mu_2$ can be written as follows:

$$m_{t+1} = \begin{cases} m_t & \text{if } \theta_{t+1} \in [\theta_{min}, \theta_{max}] \text{ and } t \notin \{t_1, t_2\} \\ (1-m_t) & \text{if } \theta_{t+1} \in [\theta_{min}, \theta_{max}] \text{ and } t \in \{t_1, t_2\} \\ 0 & \text{if } \theta_{t+1} \ge \theta_{max} \\ 1 & \text{if } \theta_{t+1} \le \theta_{min} \end{cases} \quad (22)$$

We then define $W_t^{WHs} = \{t, \dots, T\}^2 \cup \{t, \dots, T\} \cup \{\emptyset\}$ the space of at most 2 controlled switches added to a nominal trajectory.

As in Section 3.2, we define $\mathcal{X}_t^{WHs} = S^{WHs} \times W_t^{WHs}$ the controlled state space. From now on, we consider that $\forall x \in \mathcal{X}_t^{WHs}$, we can retrieve the two corresponding sequences $\{m_t(x)\}$ and $\{\theta_t(x)\}$ using (20) and (22). We define the local cost $c$ as the number of added switches:

$$c: \mathcal{X}^{WHs} \times S^{WHs} \times W_t^{WHs} \rightarrow \mathbb{R} \\ x, s, \text{switches} \quad \mapsto \#\text{switches} \quad (23)$$

We consider that all WHs have the same power consumption of 2.2kW and define the power (in kW) of one as: $P^{WHs}: \mathcal{X}_t^{WHs} \rightarrow \mathbb{R}^T$, such that $\forall t' \in \{t, \dots, T\}$:

$$P_{t'}^{\text{WHs}}: \mathcal{X}_t^{\text{WHs}} \rightarrow \mathbb{R} \\ x \quad \mapsto 2.2 \times m_{t'}(x) \quad (24)$$

## 3.4 Results

We consider a population of $N^{WHs} = 1000$ WHs with their drains sampled in $D_{test}^{WHs}$ and a population of $N^{EVs} = 1000$ EVs taken in $D_{test}^{EVs}$. The proportions are then $\alpha^{EVs} = \alpha^{WHs} = 0.5$.

Let $t \in \{1, \dots, T\}$, a step of the MPC algorithm, let $\{x^i\} \in (\mathcal{X}_t^{WHs})^{N^{WHs}}$ and

<figure><img src="image_3.png" /><figcaption>Fig. 2: MPC consumptions compared to the nominal behavior, with tracking constraints in (b) and with consumption cap constraints in (c).</figcaption></figure>
$\{y^j\} \in (\chi_t^{\text{EVs}})^{N^{\text{EVs}}}$, the controlled state of both WHs and EVs. We can define the total power consumption $P^{\text{Tot}}$ as follows:

$$ \forall t' \ge t, \quad P_{t'}^{\text{Tot}} = \sum_{i=1}^{N^{\text{WHs}}} P_{t'}^{\text{WHs}}(x^i) + \sum_{j=1}^{N^{\text{EVs}}} P_{t'}^{\text{EVs}}(y^j) \quad (25) $$

The nominal consumption is obtained using the nominal policy for the EVs and the WHs on a typical workday. We use the above equation (25) to get the power consumption during the day, as shown in Fig. 2(a).

Two different types of constraints are considered:

- The tracking constraint, using $\{r_{t'}\}_{t' \le t' \le T} \in \mathbb{R}^T$ the consumption signal to track in kW. We impose $\forall t' \in \{t, \dots, T\}, P_{t'}^{\text{Tot}} = r_{t'}$. Using it, Fig. 2(b) shows that we are able to move the consumption peak and reduce it.

- The consumption cap constraint, using $\{u_{t'}\}_{t' \le t' \le T} \in \mathbb{R}^T$ an upper bound on the consumption. We impose $\forall t' \in \{t, \dots, T\}, P_{t'}^{\text{Tot}} \le r_{t'}$. The result is shown in Fig. 2(c), with a total consumption limited to 0.7MW.

To get better results in MPC, we allow ourselves up to 8 switches per WH but still achieve the results with an average of 2 switches per WH during the day.

# 4 Conclusion and Perspectives

In this work, we have proposed a novel extension of the Moment Constrained Optimal Transport for Control (MCOT-C) framework to accommodate heterogeneous populations of controllable electrical loads. Our formulation allows for the coordination of multiple agent classes, each characterized by distinct dynamics and cost structures, while satisfying global operational constraints such as consumption caps or tracking requirements. By leveraging a mean-field approach with class-wise distributions and introducing an entropic regularization, we derive a tractable dual formulation that admits efficient numerical resolution through gradient-based methods. Furthermore, we combine this approach within a Model Predictive Control (MPC) framework, enabling online adaptation to time-varying or partially revealed information. The practical relevance of our method is demonstrated through numerical experiments on realistic datasets for electric vehicles and water heaters, showing effective constraint enforcement and interpretable decentralized policies.

Several directions for future research naturally emerge from this study. First, a theoretical analysis of the robustness of the MPC strategy to modeling errors or stochastic disturbances would strengthen the practical deployment of the approach. Secondly, the extension to dynamic or time-coupled constraints (such as ramping limitations or time-of-use pricing) may be explored to address more realistic operational scenarios. Finally, the current model relies on historical data to forecast future agent behavior but does not progressively update these forecasts using the information revealed throughout the day. Integrating
learning-based methods, such as reinforcement learning, could provide a means
to estimate unknown agent models or to improve long-term control performance
under uncertainty.

References

1. R. Carmona and F. Delarue. Probabilistic Theory of Mean Field Games with Applications I-II. Springer, 2018.

2. Y. Achdou, F. J. Buera, J.-M. Lasry, P.-L. Lions, and B. Moll. Partial differential equation models in the socio-economic sciences. Philosophical Transactions of the Royal Society A: Mathematical, Physical and Engineering Sciences, 372(2028), 2014.

3. Y. Achdou and I. Lordan. Mean field games with heterogeneity. *ESAIM: Proceedings and Surveys*, 69:1–20, 2020.

4. N. Cammarella, A. Bušić, and S. Meyn, Simultaneous allocation and control of distributed energy resources via Kullback-Leibler-Quadratic optimal control, In *Proceedings of the 2020 American Control Conference (ACC)*, pages 514–520, IEEE, 2020.

5. P. Palensky and D. Dietrich, “Demand Side Management: Demand Response, Intelligent Energy Systems, and Smart Loads,” *IEEE Transactions on Industrial Informatics*, vol. 7, no. 3, pp. 381–388, 2011.

6. D. S. Callaway and I. A. Hiskens, “Achieving controllability of electric loads,” *Proceedings of the IEEE*, vol. 99, no. 1, pp. 184–199, 2011.

7. Alfonsi, Aurélien and Coyaud, Rafaël and Ehrlacher, Virginie and Lombardi, Damiano. Approximation of Optimal Transport problems with marginal moments constraints. Mathematics of Computation, American Mathematical Society, 2020. https://doi.org/10.1090/mcom/3568.

8. Elaad OpenDataset. Available online: https://platform.elaad.io/analyses/ElaadNL_opendata.php

9. J. Albouys, N. Sabouret, Y. Haradji, M. Schumann, and C. Inard. SMACH: Multi-agent Simulation of Human Activity in the Household. Springer International Publishing. 2019.

10. J.M. Lasry and P.L. Lions. *Mean Field games*. Japanese Journal of Mathematics, 2007.

11. T. Le Corre, A. Bušić, S. Meyn, Moment Constrained Optimal Transport for Control Applications. https://arxiv.org/abs/2208.01958

12. T. Le Corre, J. Cardinal, A. Bušić, Moment Constrained Optimal Transport for Thermostatically Controlled Loads, 2025, https://arxiv.org/abs/2508.20059

13. B. M. Moreno, M. Brégère, P. Gaillard and N. Oudjane. (Online) Convex Optimization for Demand-Side Management: Application to Thermostatically Controlled Loads. HAL, 2023.

14. N. Cammarella, A. Bušić, and S. Meyn. *Kullback-Leibler-Quadratic Optimal Control*. SIAM Journal on Control and Optimization, 61(5):3234-3258, October 2023. 10.1137/21M1433654.