Cooperative multi-agent reinforcement learning forhigh-dimensional nonequilibrium control

Shriram Chennakesavalu
Department of Chemistry
Stanford University
shriramc@stanford.edu

Grant M. Rotskoff
Department of Chemistry
Stanford University
rotskoff@stanford.edu

Abstract
Experimental advances enabling high-resolution external control create new opportunities to produce materials with exotic properties. In this work, we investigate
how a multi-agent reinforcement learning approach can be used to design external
control protocols for self-assembly. We find that a fully decentralized approach
performs remarkably well even with a "coarse" level of external control. More
importantly, we see that a partially decentralized approach, where we include
information about the local environment allows us to better control our system
towards some target distribution. We explain this by analyzing our approach asa partially-observed Markov decision process. With a partially decentralized approach, the agent is able to act more presciently, both by preventing the formationof undesirable structures and by better stabilizing target structures as compared toa fully decentralized approach.

1

Context and related work

Designing nanoscale structures that are tuned to have specific material properties or dynamical
behavior is a longstanding goal in the molecular sciences [1–4]. While advances in nanofabrication
allow for increasingly intricate design, these approaches are often limited to specific materials
and require homogeneity of the molecular components. Many nanoscale systems, however, “selfassemble” in suitable conditions or in the presence of appropriate external driving forces [5–7].
Exploiting self-assembly to manufacture nanoscale components presents its own challenges, as thesystem must be designed to spontaneously and rapidly form a given target structure. The starting
materials can become kinetically trapped, forming unfavorable structures if left to self-assemble
without external control [7, 8]. Here we examine the problem of guiding self-assembly to target
high yield of the target structure via external nonequilibrium driving forces, which we refer to as a“protocol”. Reinforcement learning (RL) offers a promising toolkit for computing optimal protocols,
but there has been little systematic investigation of approaches based on RL in this context.
Here, we consider two minimal models of molecular self-assembly in which particles evolve according
to a stochastic, nonequilibrium dynamics. In both models, clusters of particles form spontaneously
under appropriate external conditions, and we seek to design protocols that optimize the size ofparticle clusters to a given target. The external driving forces (e.g., temperature, light intensity) that
we consider can be modulated as a function of both space and time (Fig. 1). In the systems considered
here, the protocol is consequently high-dimensional because the external field is modulated on a gridwith high spatial resolution. Because of the exponentially large action space, it is natural to employ
multi-agent reinforcement learning (MARL) in this context. Briefly, MARL extends the RL paradigmof an agent interacting with an environment to one in which multiple agents interact with the sameenvironment concurrently (and possibly with each other).
Fourth Workshop on Machine Learning and the Physical Sciences (NeurIPS 2021).

Several works have investigated reinforcement learning to control condensed phase dynamics, mostly
in equilibrium settings [9–12]. These approaches modulate the inter-particle interaction potentials
between constituents—fundamentally altering the constituent material—rather than externally modulating them as we do here. Recently, Ref. [13] investigated controlling dynamics of a nonequilibrium
active matter system with a low-dimensional protocol using standard RL techniques. Multi-agent
reinforcement learning has a substantial literature [14, 15]; strategies that treat each cooperative agent
independently are among the most relevant to the current study, including [16–18].
Main contributions:• We argue that within the framework of a partially observed Markov Decision Process,multi-agent reinforcement learning with improved state estimates lead to improved value.
• We demonstrate the effectiveness of incorporating local state and/or reward information
in a MARL approach for controlling self-assembly in a system of active, nonequilibrium
colloidal particles and equilibrium Lennard-Jones particles ungoing thermal annealing.

2

Multi-Agent Reinforcement Learning

We consider a partially observed Markov decision process (POMDP), cf. Ref. [19], specified asa tuple consisting of the state space, the action space, the space of observations, the observablefunction, the reward, the state-action transition operator, the discount factor, and the belief update:(S, A, Z, O, R, T, γ, b0 ). As depicted in Fig. 1, we consider an external control function that can
independently tune the strength of an external field on a specific region. We want to dynamicallyassemble clusters with a user-specified cluster size distribution ρ∗ and we takeC(x) = DKL (ρ̂(h(x))kρ∗ ).
(1)
as our cost (or inverse reward) function. We discuss in detail how we compute this reward in
Computational approach.
The function h : Rd → Nn takes a local particle configuration and simply counts the number ofclusters with k members for each k < n; this constitutes the observable O. The empirical distributionof cluster sizes is denoted ρ̂. Due to this mapping, the state information is not perfectly resolved by
the agent, but rather the input information is partially observed. We decompose the system into afixed number of independent regions, each of which is controlled by an individual agent. Because
each agent shares the same goal, the optimization is “fully cooperative” [14]. As a result, all agents
share a common cost function (1) and Q function, allowing for high data efficiency as the experienceof each agent is shared with all other agents.
This algorithm follows the MARL paradigm of centralized training but fully decentralized execution.
However, fully decentralized execution poses a limitation as spatially proximal information is crucial
for the self-assembly of clusters. Because of the high spatial resolution of control, the presence orabsence of particles and clusters in neighboring grids will no doubt influence the cluster formation
for a grid. Incorporating more local information into the observation of the current state leads toimproved belief about the current state, which we summarize in the following proposition:
Proposition 1 (Low entropy belief improves expected reward). Consider a POMDP given by(S, A, Z, O, R, T, γ, b0 ) and the POMDP augmented with local state information in which theobservation function is extended to Õ : ⊗K
i=1 Si → Z leading to a distinct evolution of the belief
probabilities, denoted b̃. Let s̄ denote a state of the system with a unique optimal action ā. If Õ issuch that b(s̄) = δs,s̄ for all states s̄, i.e. it perfectly resolves the state, then V ∗ (b̃) ≥ V ∗ (b) for allO=
6 Õ.
To prove this proposition, we simply write the Bellman optimality equation assuming the current
state is s̄, for which we know ā attains the maximal value, and letting τ denote the transition operator
on the belief vector induced by T ,
"
#
X
X
∗
∗
V (b̃) = max
R(s, a)b̃(s) + γp(z̃|b̃, a)V (τ )
a

s0

s

"
= max R(s̄, a) + γ

#
X

a

s0

2

0

∗

0

T (s |s̄, a)V (s ) ,(2)which is equivalent to the fully-observed MDP value function and z̃ = Õ(s0 ). If the observation
increases the uncertainty of the state relative to the perfectly resolved belief vector b̃, then there exists
at least one state s 6= s̄ that has belief ps > 0. In other words,
"
#
X
∗
∗
V (b) = max ps R(s, a) + (1 − ps )R(s̄, a) + γp(z̃|b̃, a)V (τ ) ≤ V ∗ (b̃),
(3)
a

s0

because the expected reward for any state s 6= s̄ is suboptimal by assumption.
Computational approach : We utilize a fully cooperative multi-agent Clipped Deep Q Learning
approach [20, 21] to learn optimized protocols that minimize C. For a given region consisting ofparticle configurations x, we define the state s = h(x) and the cost as C(x). Additionally, we consider
a discrete action space A, representing the strength of the externally modulating field on that region.
Proposition 1 suggests that incorporating additional local information should benefit the optimization,
so long as that information improves the accuracy of the belief about the current state. To do so, we
consider a hierarchy of local region sizes: a plaquette consisting of the current region and its nearest
neighbors, a 3 × 3 grid, a 5 × 5 grid and the global system. Given particle configurations in the local
regions, x0 , we define the state of the local region s0 = h(x0 ) and the cost of the local region as C(x0 ).
With this in mind, we extend our approach to include centralized training but partially decentralized
execution. This corresponds to sharing information among the agents, as we include information
about the state and/or cost of local regions. When including state information about the local region,
we consider the joint state between a region s and its surrounding local region s0 . Additionally, when
including cost information about the surrounding region, we consider the average cost between C(x)and C(x0 ). We indicate whether or not we included surrounding state and/or cost information with asubscript S for cooperative state and a subscript C for cooperative cost.

3

Active Colloids

Fig. 1 summarizes numerical results for a 2D colloidal particle system driven by externally controlled
light [22]. In this system, when the Péclet number is sufficiently large (i.e. when the light-induced
activity is high relative to the temperature), the system exhibits a nonequilibrium phase separation,
called motility induced phase separation (MIPS) [23].
To explore the limits of an activity-inducing external control, we sought to maintain a steady state
distribution consisting of clusters of particles much smaller than the macroscopic aggregate that forms
when there is constant activity. We specified a target distribution ρ∗ of cluster sizes using a Gaussian
target distribution with µ = 20, σ 2 = 9 and considered a 48 × 48 grid of control. Consistent with
Prop. 1 including information of surrounding grids improves the performance of the optimal protocol.
The best performing protocol, 3 × 3CS , includes cluster size information and cost information ina local 3 × 3 region. Surprisingly, including just the cost information of this local region, 3 × 3C ,actually hindered the performance compared to the fully decentralized approach. On the other hand,for the GlobalC case, when we included information about only the cost of the entire global system,
we see that the performance does remarkably well. This underscores the importance of carefully
selecting relevant centralized or semi-centralized information to improve performance. Acting in apartially decentralized manner can not only incur an additional computation cost but can also resultin a decrease in performance.
Including surrounding information in the best performing approach was important in two cases: when
the current region was empty or when it contained a cluster that was at or near the mean of the target
distribution. In both these cases, in a fully decentralized setting, the trained agent always selects anaction with the highest activity. However, in the partially decentralized setting, the agent enforces alower activity when the surrounding region contains clusters that are above the mean of the target
distribution, preventing the formation of a cluster far larger than the mean target size.

4

Thermal Annealing

Thermal annealing is widely used to improve the yield of nanoscopic materials [24, 25]. Annealing islimited as a mechanism for control because there is essentially only one parameter that can be tuned,3

2.5

DKL(ρ̂(τ )kρ∗)

3.0

7

2.0

|µρ̂(τ ) − µρ∗ |6Fully Decentralized
GlobalC
Nearest NeighborsCS
3 × 3CS5 × 5CS3 × 3C

1.5

5

1.0

4

0.5

3

0.0
0

500

1000

1500

2000

2500

3000

0

Time (s)

500

1000

1500

2000

2500

3000

Time (s)
1.4
1.2
1.0
0.8
0.6
0.4
0.2
0.0
1.4
1.2
1.0
0.8
0.6
0.4
0.2
0.0

Figure 1: Active Colloids. In the plots above, both the KL divergence from the instantaneouscluster size distribution to the target distribution and the absolute deviation of the mean instantaneouscluster size distribution µρ̂(τ ) from the mean target size µρ? are shown as a function of time for the
optimized protocols of various partially decentralized approaches and a fully decentralized approach(the protocols are fully trained in these curves). In the illustrations below, we compare the predictedaction in the fully decentralized case (no gray shading) and the best performing partially decentralized
case, 3 × 3RS (with gray shading). Higher light-induced activities (red) promote cluster formation.
the rate at which the temperature is decreased. We examine an alternative paradigm that exercises
more localized control with measurement-guided feedback to design an annealing schedule. Rather
than globally tuning a temperature, we locally update the temperature on a grid, see Fig. 2. As above,
we fixed a target cluster size distribution, chosen to be a Gamma distribution with variance σ 2 = 1
and a mean µ = 4 and considered a 15x15 grid of control. Here, the size and density of the systemis far lower and we sought to form clusters of a much smaller size compared to the active system.
This allowed us to define individual grids that were only slightly larger than an individual particle,
providing much finer levels of control. Note that the interaction potential for these particles is distinct
from the one that we use to model the active system and there is no nonequilibrium dynamics. The
fully decentralized and top performing partially decentralized approaches differed in their behavior
when the cluster size was around the mean. In the fully decentralized case, the agent was far moreconservative when a cluster size was of size 4 or greater, and would impose a high temperature
breaking apart the cluster. On the other hand, in one of the top performing partially decentralized
approaches, Nearest NeighborsCS , the agent allowed cluster sizes between 3 and 5 to persist by
ensuring a lower temperature, given that there were no clusters in the vicinity.

Broader Impact
In this work, we investigated how a partially decentralized reinforcement learning could be used
to learn protocols that externally drive a system towards some non-equilibrium steady-state. Thisapproach can be used as a framework for learning protocols for self-assembly in more complex
systems. For example, in the design of suprastructures, each agent—or a local grouping of agents—
can be used to design individual substructures, which then self-assemble to form larger structures.
4

1.6

1.0

1.4

DKL(ρ̂(τ )kρ∗)|µρ̂(τ ) − µρ∗ |

1.2

0.4

0.4

0.8

Fully Decentralized
GlobalCS
Nearest NeighborsCS
3 × 3CS5 × 5C

1.0

0.8

0.6

0.6

0.2

0.2

0.0
0

500

1000

1500

2000

2500

3000

0

Time (s)

500

1000

1500

2000

2500

3000

Time (s)
1.0
0.8
0.6
0.4
0.2

1.0
0.8
0.6
0.4
0.2

Figure 2: Thermal Annealing. See Fig. 1 caption for descriptions of plots. In the cartoons
below, we compare the predicted action in the fully decentralized case (no gray shading) and
Nearest NeighborsCS (with gray shading). Lower temperatures (more blue) promote cluster formation.
At present, the key limitation is that the paradigm is most naturally suited to control with spatial
resolution. Ultimately, the ability to design protocols for self-assembly of nanoscopic structures, canhave wide-ranging impacts across a number of disciplines ranging from medicine to materials science.
Potential negative impacts could, in principle, arise from the design of harmful or toxic materials.
Data and Code Availability: The data that support the findings of this study are available fromthe corresponding author upon reasonable request. Our code is available at https://github.com/
rotskoff-group/marl-design.

Checklist
1. For all authors...
(a) Do the main claims made in the abstract and introduction accurately reflect the paper’s
contributions and scope? [Yes](b) Did you describe the limitations of your work? [Yes](c) Did you discuss any potential negative societal impacts of your work? [Yes](d) Have you read the ethics review guidelines and ensured that your paper conforms tothem? [Yes]2. If you are including theoretical results...
(a) Did you state the full set of assumptions of all theoretical results? [Yes](b) Did you include complete proofs of all theoretical results? [Yes]3. If you ran experiments...
(a) Did you include the code, data, and instructions needed to reproduce the main experimental results (either in the supplemental material or as a URL)? [Yes]5

(b) Did you specify all the training details (e.g., data splits, hyperparameters, how they
were chosen)? [Yes](c) Did you report error bars (e.g., with respect to the random seed after running experiments multiple times)? [Yes](d) Did you include the total amount of compute and the type of resources used (e.g., typeof GPUs, internal cluster, or cloud provider)? [No]4. If you are using existing assets (e.g., code, data, models) or curating/releasing new assets...
(a) If your work uses existing assets, did you cite the creators? [N/A](b) Did you mention the license of the assets? [N/A](c) Did you include any new assets either in the supplemental material or as a URL? [N/A](d) Did you discuss whether and how consent was obtained from people whose data you’reusing/curating? [N/A](e) Did you discuss whether the data you are using/curating contains personally identifiable
information or offensive content? [N/A]
5. If you used crowdsourcing or conducted research with human subjects...
(a) Did you include the full text of instructions given to participants and screenshots, if
applicable? [N/A](b) Did you describe any potential participant risks, with links to Institutional Review
Board (IRB) approvals, if applicable? [N/A](c) Did you include the estimated hourly wage paid to participants and the total amountspent on participant compensation? [N/A]

References
1

Y. Yin and A. P. Alivisatos, “Colloidal nanocrystal synthesis and the organic–inorganic interface”,
Nature 437, 664–670 (2005).
2
Y. Ma and A. L. Ferguson, “Inverse design of self-assembling colloidal crystals with omnidirectional
photonic bandgaps”, Soft Matter 15, 8808–8826 (2019).
3
K. R. Gadelrab, A. F. Hannon, C. A. Ross, and A. Alexander-Katz, “Inverting the design path forself-assembled block copolymers”, Molecular Systems Design & Engineering 2, 539–548 (2017).
4
H. Ronellenfitsch, N. Stoop, J. Yu, A. Forrow, and J. Dunkel, “Inverse design of discrete mechanicalmetamaterials”, Physical Review Materials 3, 095201 (2019).
5
J. C. Cameron, S. C. Wilson, S. L. Bernstein, and C. A. Kerfeld, “Biogenesis of a Bacterial
Organelle: The Carboxysome Assembly Pathway”, Cell 155, 1131–1140 (2013).
6
C. Sigl et al., “Programmable icosahedral shell system for virus trapping”, Nature Materials, 1–9(2021).
7
G. M. Rotskoff and P. L. Geissler, “Robust nonequilibrium pathways to microcompartment assembly”, Proceedings of the National Academy of Sciences 112, 201802499 (2018).
8
S. Whitelam and R. L. Jack, “The Statistical Mechanics of Dynamic Pathways to Self-Assembly”,
Annual Review of Physical Chemistry 66, 143–163 (2015).
9
M. C. Rechtsman, F. H. Stillinger, and S. Torquato, “Optimized Interactions for Targeted SelfAssembly: Application to a Honeycomb Lattice”, Physical Review Letters 95, 228301 (2005).
10
D. Chen, G. Zhang, and S. Torquato, “Inverse design of colloidal crystals via optimized patchy
interactions”, The Journal of Physical Chemistry B 122, 8462–8468 (2018).
11
F. Romano, J. Russo, L. Kroc, and P. Šulc, “Designing Patchy Interactions to Self-Assemble
Arbitrary Structures”, Physical Review Letters 125, 118003 (2020).
12
Z. Ma, E. Lomba, and S. Torquato, “Optimized large hyperuniform binary colloidal suspensions in
two dimensions”, Physical Review Letters 125, 068002 (2020).
13
M. J. Falk, V. Alizadehyazdi, H. Jaeger, and A. Murugan, “Learning to Control Active Matter”,
arXiv:2105.04641 [cond-mat] (2021).
14
M. Tan, “Multi-agent reinforcement learning: Independent versus cooperative agents”, in Machinelearning, proceedings of the tenth international conference, university of massachusetts, amherst,
MA, USA, june 27-29, 1993, edited by P. E. Utgoff (1993), pp. 330–337.
6

15

L. Busoniu, R. Babuska, and B. De Schutter, “A Comprehensive Survey of Multiagent Reinforcement Learning”, IEEE Transactions on Systems, Man, and Cybernetics, Part C (Applications andReviews) 38, 156–172 (2008).
16
T. Rashid, M. Samvelyan, C. S. de Witt, G. Farquhar, J. N. Foerster, and S. Whiteson, “QMIX:Monotonic value function factorisation for deep multi-agent reinforcement learning”, in Proceedings of the 35th international conference on machine learning, ICML 2018, stockholmsmässan,
stockholm, sweden, july 10-15, 2018, Vol. 80, edited by J. G. Dy and A. Krause, Proceedings ofmachine learning research (2018), pp. 4292–4301.
17
K. Son, D. Kim, W. J. Kang, D. Hostallero, and Y. Yi, “QTRAN: Learning to factorize withtransformation for cooperative multi-agent reinforcement learning”, in Proceedings of the 36th
international conference on machine learning, ICML 2019, 9-15 june 2019, long beach, california,
USA, Vol. 97, edited by K. Chaudhuri and R. Salakhutdinov, Proceedings of machine learning
research (2019), pp. 5887–5896.
18
P. Sunehag et al., “Value-decomposition networks for cooperative multi-agent learning based onteam reward”, in Proceedings of the 17th international conference on autonomous agents andMultiAgent systems, AAMAS 2018, stockholm, sweden, july 10-15, 2018, edited by E. André,
S. Koenig, M. Dastani, and G. Sukthankar (2018), pp. 2085–2087.
19
N. Roy, G. Gordon, and S. Thrun, “Finding Approximate POMDP solutions Through Belief
Compression”, Journal of Artificial Intelligence Research 23, 1–40 (2005).
20
V. Mnih, K. Kavukcuoglu, D. Silver, A. Graves, I. Antonoglou, D. Wierstra, and M. Riedmiller,“Playing atari with deep reinforcement learning”, in NIPS deep learning workshop (2013).
21
S. Fujimoto, H. van Hoof, and D. Meger, “Addressing function approximation error in actor-critic
methods”, in Proceedings of the 35th international conference on machine learning, Vol. 80, edited
by J. Dy and A. Krause, Proceedings of machine learning research (July 2018), pp. 1587–1596.
22
J. Palacci, S. Sacanna, A. P. Steinberg, D. J. Pine, and P. M. Chaikin, “Living Crystals of LightActivated Colloidal Surfers”, Science 339, 936–940 (2013).
23
M. E. Cates and J. Tailleur, “Motility-Induced Phase Separation”, Annual Review of Condensed
Matter Physics 6, 219–244 (2015).
24
G. Makrides, B. Zinsser, A. Phinikarides, M. Schubert, and G. E. Georghiou, “Temperature andthermal annealing effects on different photovoltaic technologies”, Renewable Energy 43, 407–417(2012).
25
S. Dey et al., “DNA origami”, Nature Reviews Methods Primers 1, 1–24 (2021).

7