SwarmRL: Building the Future of Smart Active Systems
Samuel Tovey Christoph Lohrmann Tobias Merkt David Zimmer Konstantin Nikolaou Simon Koppenhöfer Anna Bushmakina Jonas Scheunemann Christian Holm

Supplementary Material for 
 “Revisiting Relevance Feedback for 
 CLIP-based Interactive Image Retrieval”
    Ryoya Nara1Yu-Chieh Linr2 Yuji Nozawa2 Youyang Ng2 Goh Itoh2 Osamu Torii2 Yusuke Matsui1

    April 29, 2024
=========================================================================================================

[2]These authors contributed equally

This work introduces , a Python package designed to study intelligent active particles.
 provides an easy-to-use interface for developing models to control microscopic colloids using classical control and deep reinforcement learning approaches.
These models may be deployed in simulations or real-world environments under a common framework.
We explain the structure of the software and its key features and demonstrate how it can be used to accelerate research.
With , we aim to streamline research into micro-robotic control while bridging the gap between experimental and simulation-driven sciences.
 is available open-source on GitHub at <https://github.com/SwarmRL/SwarmRL>.





§ INTRODUCTION


Mastering control of microrobots at the microscopic scale has the potential for insights and the development of new technologies capable of changing the world.
Whether it be an improved understanding of bacterial navigation strategies or direct control over microscopic robots, numerous fields, including construction <cit.>, plant pollination and ecosystem defense <cit.>, and search and rescue <cit.>, will be enhanced by advancements in micro-scale control research.
One field in particular that will see significant changes is medicine, where it is expected that micro robotic agents will be capable of advanced treatment strategies including targeted cancer therapies <cit.>, assisted and drug delivery <cit.>, assisted fertilisation <cit.>, amongst many others <cit.>.
The two most significant challenges in realising this future are developing the tools and materials to build such capable microscopic agents and programming them to achieve complex tasks with minimal observed input.
These agents should not only be capable of performing complex actions such as swimming, pushing, and perhaps, communications, but they should do so without constant input from a supervisor.
The first of these challenges has been, and is being, extensively investigated.
Currently, there are many approaches for designing mobile micro robots <cit.>.
Initial approaches included the light-driven Janus particles <cit.> and coils <cit.> and magnetically driven devices taking on numerous shapes <cit.>.
Recent approaches have also included using living cells to produce motion on this microscopic scale, the Xenobots being a prime example <cit.>.
However, the control of these agents, once built and deployed, is a complicated problem.
Due to the size of these agents, directly installing computational processing power onto them is, thus far, intractable.
Therefore, simplistic algorithms that aim to encode reactionary behaviour into the agents are often constructed, something that would require limited on-board processing.
Examples of this include the use of microrobots to perform chemotaxis <cit.> by changing their shapes upon exposure to varying fields, object manipulation <cit.> through the use of an external magnetic field, and reproducing swarming behaviour by applying simple rotations and translations to individual agents by a set of rules triggered by changes in local environment <cit.>.
More recent approaches have utilised reinforcement learning approaches to further push in the direction of learned control strategies.
<cit.> utilised a Q-learning algorithm to learn swimming strategies in gated microswimmers, <cit.> used an actor-critic framework to learn competitive capture-evasion strategies from hydrodynamic cues, and <cit.> investigated actor-critic learned navigation strategies under the influence of stochastic environments.
In all cases, demonstrating under what conditions microrobots can perform complex tasks in various environments brings us a step closer to realising the full potential of this technology.

It is clear that research in the direction of micro-robotic control is both promising and excitingly multi-disciplinary; however, due to the nature of the technology required, whether it be complex experimental equipment, deep knowledge of state-of-the-art machine learning algorithms and their implementation, or the expertise to construct physically realistic simulations, the entry barrier can be high for specialists who focus on only one of these areas.
For this reason, we have developed the  Python package to combine each control workflow element under a common framework and accelerate algorithm development and deployment.
 enables researchers to apply control strategies to particles in simulation or experiments.
These strategies can be built from sets of rules that are rigidly programmed and utilise the environment of the particles, or they can be trained using the actor-critic reinforcement learning approach, offering complete flexibility in how the agents are controlled, what actions they are capable of, and what tasks they must perform.
Beyond the customizability, wherever possible, it is built on top of the JAX <cit.> ecosystem to maximise performance and is suited for deployment on large distributed compute clusters.
This paper aims to introduce the  software and demonstrate how it can be applied to perform frontier research.
We begin with an overview of the theory underpinning the software, explicitly discussing microscale active matter, simulating it in a physically realistic manner, and the reinforcement learning we have built to control it.
Afterwards, the architecture of the software is described, detailing how each piece of the software fits together and what can be done with them.
Finally, we discuss some of the unique features of , including performance and visualisation capabilities, before concluding with an outlook of the project.
With , we hope to enable scientists from diverse backgrounds to contribute to and develop the field of micro-robotics research and realise the potential of this technology.


§ THEORY


 harnesses tools from a range of fields, the largest of which are active matter simulation and reinforcement learning.
In order to understand and better utilize the software infrastructure, it is important to also have an understanding of the theory it implements. 
In this section, we introduce the important theoretical concepts on which  is built, before exploring the architecture of the package.


 §.§ Active Matter




  §.§.§ Biological and artificial active systems

Active matter refers to systems that consume energy from their surroundings and are thus internally driven out of equilibrium <cit.>.
All biological and human-made machines fall under this term. 
However, in the context of this work, we focus on the branch of active systems that consists of small active particles that use energy to perform persistent, directed motion.
This much more narrow definition covers plenty of examples in biology and engineering.
Many bacteria like Escherichia coli <cit.>, archaea like Thermoplasma volcanium <cit.>  and small eukaryotes like Chlamydomonas reinhardtii <cit.> use molecular motors and organelles like archaella, flagella or pili to self-propel.
Artificial microswimmers can be made from colloidal particles that use, e.g., catalytic reactions <cit.>, self-generated temperature gradients <cit.> or magnetic fields <cit.> to generate propulsion.
In some more rare cases, the active particles can self-propel and self-steer, i.e., use torque to actively change their direction.
For bacteria, tumbling <cit.> is a well-known mechanism for change of direction.
Some artificial microswimmers can also be steered by, e.g., change in laser focus <cit.> or magnetic field direction <cit.>.
In all these examples, the active particle is of micrometer size and suspended in a liquid, making translational and rotational Brownian motion a non-negligible factor that competes with the deterministic active motion.



  §.§.§ Simulation

Active matter systems with self-propelled small particles are usually modelled within the framework of stochastic dynamics of the individual constituents. 
We use the two- or three dimensional overdamped Langevin equations










    _i(t) = γ_t^-1[_i(t) _i(t) + F_i(_i, _j) ] + √(2  T γ_t^-1)^t_i(t),


    _i(t) = [γ_r^-1[_i(t) _i(t) + M(_i, _j) ] + √(2  T γ_r^-1)^r_i(t)] _i(t),

where _i denotes the position of particle i, γ_t(r) the translational (rotational) friction coefficient,  the active driving force that models self-propulsion,  the unit vector representing the direction of the particle, F_i(_i, _j) a conservative force from interactions between particle i and its environment as well as from interactions with other particles j,  the Boltzmann constant, T the absolute temperature, ^t,(r) the translational (rotational) noise with ^t,(r)_i(t) = 0 and ^t,(r)_i(t)⊗^t,(r)_j(t') = δ_ijδ(t-t')1, where · denotes an ensemble average (expectation value) and 1 the unity matrix,  an active steering torque,  a unit vector perpendicular to  that sets the steering direction and M_i(_i, _j) an interaction torque analogous to F_i.
We stress the time dependence of  and  because  is designed to focus on control and decision-making on the level of single particles.
Collective behavior and task fulfillment are to be achieved by the particles' actions and what they can control rather than by external fields acting on all particles.




 §.§ Reinforcement Learning

As a sub-field of ML, RL is concerned with how intelligent agents can take actions in an environment to maximize a cumulative reward. 
It allows these agents to learn optimal protocols to achieve a given task or a set of tasks without explicitly telling them how. 
This is oftentimes useful when a task is too complex for a simple algorithm or control system.
It can also be used when the focus of the investigation is the emergence of a strategy for a problem.
RL is, however, data-inefficient and is thus used mostly for problems where data exists abundantly and where the action space is relatively small such as autonomous driving <cit.>, video games <cit.>, and robotics <cit.>. 
It reached great popularity in 2016 when Deepmind's AlphaGo won against the top Go player Lee Sedol <cit.>.
When using it for specific tasks, however, RL can achieve remarkable results. 
In 2022, it solved a 50-year-old open question in mathematics how to most efficiently use matrix multiplication <cit.>; this discovery likely improves the efficiency of machine learning significantly due to its widespread use of matrix multiplication. 
It is further increasingly used in Physics research for sophisticated tasks such as quantum error correction <cit.> and the magnetic control of tokamak plasmas <cit.>. 
This section provides an introduction to RL and explains the training approach currently implemented in SwarmRL, namely, deep Actor-Critic (AC).



  §.§.§ Actor-Critic Reinforcement Learning

The RL training paradigm is concerned with iteratively improving an agent's decision-making such that it can achieve a pre-defined task as efficiently as possible.
In practice, this is performed by placing agents into an environment with which they can interact by performing actions, a_t.
To decide on such an action, the agent receives a state description, s_t, that describes the environment with sufficient information.
The model used to decide on the action is referred to as the policy, π_θ: s_t→ a_t, often taken to be a neural network with parameters, θ that is then referred to as an .
These actors typically produce a probability distribution, which can then be sampled to select an action given the state.
After an action is taken, a new state is produced along with a reward, quantifying the agent's progress towards achieving its task, a process outlined in Figure <ref>.

Fundamentally, the goal of training in RL is to find the optimal policy that, throughout a trajectory, τ, maximises a cost function, J(π_θ), i.e.,

    π^* = _π J(π_θ) = _π∫_τP(τ | π_θ)· R(τ) = _π⟨ R(τ) ⟩_τ∼π_θ,

where R(τ) is the cumulative reward computed from a trajectory, P(τ | π_θ) is the probability of a trajectory given the current policy, and the angled brackets denote an average over many trajectories sampled under the π_θ policy.
This is achieved by encouraging the policy to increase the probability of selecting actions that returned good rewards and discouraging those that didn't.
In practice, training is done throughout discrete-time episodes where state, policy probabilities, and reward data are collected and used to update the policy via the gradient ascent algorithm

    θ^' = θ + η·∇_θJ(π_θ),

with learning rate η, and

    ∇_θJ = ⟨∑_t^T∇_θlogπ_θ(a_t|s_t)· r_t⟩_τ∼π_θ

with time step, t, and instantaneous reward, r_t.
Typically, to ensure convergence in the long-term rewards and to encourage near-term over long-term progress, an expected returns function, G_t, will be used in the form

    G_t = ∑_t^' = t^Tγ^t^' - tr_t^',

where T is the final time step of the episode, and γ is a discount factor set as a hyperparameter during training influencing the importance of late-time rewards.
While this approach can be sufficient in training agents, it can break down in stochastic environments and high-dimensional spaces <cit.>.
To address this, <cit.> showed that introducing a trainable baseline or value function, V^π to compute an advantage function

    A_t^π = G_t - V^π_t,

can result in the stabilization of the expected returns and improved convergence time in the training, resulting in the reformulation of Equation <ref> to

    ∇_θJ = ⟨∑_t^T∇_θlogπ_θ(a_t|s_t)· A^π_t⟩_τ∼π_θ.

When this value function is implemented as a neural network, it is often referred to as a , thus actor-critic reinforcement learning.
The critic's role is to determine the value of the state in which the agent currently finds itself.
While several definitions exist, in , the value is defined as the theoretical returns possible if one starts in a state, s_t and follows policy π_θ. 
It improves this prediction by training on the true expected returns, Equation <ref>, computed during episodes through a desired regression algorithm.
Therefore, in Equation <ref>, the advantage function guides training by encouraging selected actions that are better than or equal to those expected from the value function and discouraging those considered below what is theoretically possible.
For this reason, AC RL training is often considered adversarial in nature, with the actor learning to beat the critic and the critic, in turn, learning the correct value of the encountered states.
The training procedure outlined here is referred to as Vanilla Policy Gradient due to it being the simplest form of actor-critic training.
More complex methods extend on the method discussed here, such as trust region policy optimisation <cit.> and proximal policy optimisation <cit.>, the latter of which is also implemented in .



  §.§.§ Multi-Agent Reinforcement Learning

As  specializes in controlling microscopic agents, there will be many occasions in which multiple agents must be simultaneously controlled and trained.
This is addressed by multi-agent reinforcement learning (MARL) <cit.>. 
In MARL, N autonomous agents interact with the environment and each other. 
This extension brings several new aspects to the problem, such as cooperation and competition among agents and various information structures. 
In the following, we will address these aspects briefly. 


  
Cooperative, Competitive and Mixed Setting
In a fully cooperative setting, all agents share a common reward, i.e., r^1=r^2=...=r^N = r or a team-average reward where r=1/N∑_i∈ Nr_i <cit.>. 
The latter is a generalized version of the former. 
Different reward functions for each agent allow more heterogeneity among agents <cit.>.
A fully competitive setting in MARL is typically modelled as a zero-sum Markov game with ∑_i∈ Nr^i=0 <cit.>. 
Most literature focuses on two players where one agent's reward is the other's loss, something leveraged in, for example, predator-prey style games.
A mixed setting imposes no restrictions on the goal and the relationship among agents <cit.>. 


  
Information Structures
Another critical challenge in MARL is the question who knows what during training and execution. 
This involves both the decision-making process as well as the optimization process. 
Three different forms of information structure are generally accepted.
The first is the centralized setting with a so-called central controller aggregating joint observations, actions, and rewards and allowing for centralized learning. 
The execution, on the other hand, can either be centralized or decentralized.
A centralized executed policy produces a joint action for all agents, while decentralized produces 
individually executed policies. 
The latter is the popular learning scheme of centralized-learning-decentralized-execution <cit.>.
On the other side, there is the fully decentralized setting in which each agent acts based only on local observation. 
This setting results in the so-called independent learning scheme <cit.>.
A combination of both is a decentralized setting with networked agents <cit.>. 
This allows agents to share local information with their neighbours.
The choice of setting depends on the problem at hand. 


§ ARCHITECTURE


 has been designed to lower the entry barrier for scientists to access state-of-the-art reinforcement learning algorithms efficiently.
It also offers a wide degree of customization and choice of algorithms that can be used during experiments to maximise the experienced user's flexibility.
To achieve both of these goals,  is constructed using a modular, object oriented programming structure where almost all pipeline components are written as implementations of an abstract parent that defines the interface through which the components communicate.
The Python programming language has been used throughout as it is most common amongst scientific disciplines, however, whenever possible, libraries have been used to improve performance through vectorized maps (vmaps), multi-threading, GPU use, and Just In Time (JIT) compiling <cit.>.
A coarse-grained flowchart of  is presented in Figure <ref> where each major component of the architecture is shown with its corresponding connections.
The structure of  is built around the interaction between the , whether that be simulation or real experiment, and the , the algorithms used to control the particles in the simulation.
Connecting the  to the  is the .
Force Functions are either constructed directly, or through a  in the cases where trainable  are in use.
Finally, the  holds , , and , each of which provide information and functionality to the algorithm being used.
The rest of this section is dedicated to explaining each important component of , where they fit into the different capabilities of the software, and how they can be adapted or extended for individual implementations.
It should be noted that, to reduce space and improve readability, all comments, type hints, and methods not directly related to the discussion have been removed from code samples.



 §.§ Engine


The  class of  is one of its most important components.
It represents the environment of the  and propagates the state according to the given actions.
 ships with an engine connected to the ESPResSo simulation software <cit.> to solve Equations (<ref>) and (<ref>) numerically.
Through ESPResSo, we also offer the underdamped Langevin equation and a coupling to explicit hydrodynamics with lattice Boltzmann as the solver of the Navier-Stokes equations.
 also ships with rudimentary examples for connecting it directly to an experiment. 
Direct experimental connections are often kept private out of respect to the institutions that host the experiments.
However, most implementations thus far use the TCP communications protocol to send actions and receive state information to hardware.
Future work on  aims to unify this communication procedure using the ROS ecosystem <cit.>.
In cases where users must connect their own engines, be it simulation or experiment, they inherit from the  parent class. 
They must implement the two methods outline in Code Sample <ref>.
[language=Python, label=lst:engine, caption=  parent class.]
class Engine:

    def integrate(
        self, n_slices, force_model,
    ):
        raise NotImplementedError

    def get_particle_data(self):
        raise NotImplementedError
        

The  method is used to step forward either in simulation or in an experiment.
It takes as arguments the force model, which holds the control algorithms for the particles in the experiment, see <ref>, and number of slices to be looped over, that is, the number of times the force model is called for new actions to be applied to the particles.
In a simulation, the time in-between successive time slices is filled with integrator steps, applying the chosen action at each time-step. 
In contrast, in experiment, this would be applying the actions for a fixed period between slices.
When instantiated, these parameters, time-step or time between actions, are given to the specific engine.
The second method that must be implemented is .
This method must collect information about the particles in the system and give it back to the user uniformly so that all parts of  can work with them.
In its current state,  defines colloids with properties outlined in Code Sample <ref>.
[language=Python, label=lst:colloid, caption=  dataclass.]
class Colloid:
    pos
    director
    id
    velocity
    type

Here, all information relevant to the models is assigned to a colloid class and used to predict new actions and to compute rewards.


 §.§ Agents


 was built to service a variety of simulations and RL studies.
A large component in enabling this, is the  class, which manages how agents in the simulations are controlled, what they see, and what they are tasked with doing.
The  collects the components introduced in the previous sections and represents the full set of sensing, control, and objectives of the colloids in question.
This information can then be used by  to effectively deploy these agents in the simulation environment.
Code samples <ref> and <ref> demonstrate the construction of two different agents which could both be added to a single simulation within .

[label=lst:ac-agent, caption=Actor-Critic Agent, language=Python]Name
ac_agent = srl.agents.ActorCriticAgent(
    agent_type=0,
    network=srl.networks.FlaxNetwork(...),
    task=srl.tasks.RotateRod(),
    observable=srl.observables.VisionCones(),
    actions=action_dict
)


[label=lst:class-agent, caption=Classical Agent, language=Python]Name
ac_agent = srl.agents.Lymburn(
    agent_type=1,
    homing_location=[500., 500., 500.]
)


In the above examples, the actor-critic reinforcement learning approach controls one group of particles. 
This means their policy can be trained and adapted to new environments.
The other example uses a so-called , in particular, one based on the swarming agents of <cit.>.



 §.§ Classical Control


, allows users to control the actions of the particles in the system using both RL and classical algorithms.
Classical algorithms, in this work, can be implemented without training and typically follow a defined set of rules.
Such sets of rules have often been studied in the literature to reproduce behaviour in microswimmers <cit.>.
Classical agents are implemented directly into the  child class on the software side. 
However, they must follow the same structure, namely, at each time slice, the information of all particles in the system is passed into the class and in turn, it must return an action for each particle.
In its current state,  offers classical algorithms for the published models of <cit.>, <cit.>, and <cit.>.
Future work aims at including Model Predictive Control (MPC) algorithms into the  ecosystem.


 §.§ Machine Learning Control


The machine learning side of agent control is driven by multi-agent reinforcement learning.
However, due to the complex nature of this approach, several design decision have been made to balance flexibility with limiting complexity.
The main workhorse of the machine learning driven control is the  class, which requires two methods to be implemented,  which computes the action of the agent/s under study, and , which returns action probabilities and, in the case of actor-critic approaches, also the predicted state values.


  §.§.§ Architecture

 currently utilizes the Flax neural network library <cit.> exclusively.
In order to provide an interface for complex architectures,  requires only a single Flax module to be passed to the  class.
This module must return both the action probabilities, and the predicted value function in case of actor-critic training.
This approach allows implementing very complex neural network architectures such as graph models, with the same API interface as simple dense actors and critics.
An example of two commonly used approaches is listed in Code Samples <ref> and <ref>.

[label=lst:disjoint, caption=Disjoint actor-critic, language=Python]Name
class DisjointActorCritic:

  def __call__(self, x):
    # Perform the actor forward pass.
    actor_values = Dense(12)(x)
    actor_values = relu(actor_values)
    actor_output = Dense(4)(actor_values)

    # Perform the critic forward pass.
    critic_values = Dense(12)(x)
    critic_values = relu(critic_values)
    critic_output = Dense(1)(critic_values)

        return actor_output, critic_output        


[label=lst:joint, caption=Combined actor-critic, language=Python]Name
class CombinedActorCritic:

  def __call__(self, x):
    # Perform the shared layer computation.
    shared_layer = Dense(12)(x)
    shared_layer = relu(shared_layer)

    # Compute the actor output.
    actor_output = Dense(4)(shared_layer)

    # Compute the critic output.
    critic_output = Dense(1)(shared_layer)

        return actor_output, critic_output      


These code blocks show a disjoint and combined actor critic approach bundled into the same Flax network structure.
In this way, many complex neural network architectures can be constructed under a single module including graph or transformer modules.

Apart from the architecture, several components are essential to training AC models: the policy update algorithm, the approach used in the expectation computation, and the strategy for sampling actions from the learned distribution.




  §.§.§ Actions

Actions in  are dictated by  which can be read by an  and turned into forces, torques, or other values in a simulation or experiment.
At the time of this report,  has been focused on micro-robotics with standard direction control, therefore, the class is built from a simple set of attributes outlined in Code Sample <ref>.
[language=Python, label=lst:action, caption=  dataclass.]
class Action:
  force = 0.0
  torque= np.zeros((3,))
  new_direction = None

Force is applied along the agent's director axis, torque describes the agent's rotation, and new_direction provides a sometimes simpler approach to a torque calculation.
From this template, we create an action by specifying the class's force, torque, and new direction attributes and add them to a dictionary to be passed to an agent.
For example, if the agent should be capable of translation forwards, translation backwards, clockwise, and counterclockwise rotation, the set of actions can be passed to the  as in Code Sample <ref>.
[language=Python, label=lst:actionex, caption=Example  dictionary.]
agent_action_dictionary = 
  "Forwards": Action(force=10.0),
  "Backwards": Action(force=-10.0),
  "CCW": Action(torque=np.array(0., 0., 10.))
  "CW": Action(torque=np.array(0., 0., -10.))





 §.§ Force Function


The interaction model in  is responsible for consuming data from the simulation and sending back computed actions for each of the .
 can control agents using machine learning and standard control theory approaches.
The parent class for the models requires only a single function to be called by the : , which takes a list of  objects as input and returns a list of .
This is the only occasion in which the environment communicates with the control algorithms and the communication must adhere to this very 'narrow' interface.
This design choice enables novice  users to switch out environments without adapting their RL setup and experienced users to easily implement new environments according to their needs.

Agent information is passed through as a dictionary when creating a model, demonstrated in Code Sample <ref>.
[language=Python, label=lst:model, caption=Example  instantiation for a classical and actor-critic RL agent.]
interaction_model = ForceFunction(
    agents=
        "0": ActorCriticAgent,
        "1": ClassicalAgent
    
)

In the case above, the experiment has at least two types of particles.
One is controlled using an actor-critic reinforcement learning approach while the other takes actions computed by a classical, non-trainable algorithm.
In this way,  users have complete control over the agents in their simulations and experiments.



 §.§ Tasks


As the goal of  is to be used for the study of micro-scale robotic or biological systems to learn about their behaviour and control, it is often important to measure how well the system is doing.
In RL, this is often referred to as the reward which is maximised during training.
For classical algorithms, a task can also be used as a measure of how well the approach is doing or passed into the model as feedback.
In , system performance measurement is handled by the  class.
The parent class of the task is built up of several methods that need to be implemented in the case of a custom task.
The reduced class architecture is displayed in Code Sample <ref>.
[language=Python, label=lst:task, caption=  parent class.]
class Task:
  @property
  def kill_switch(self):
      pass
  def __call__(self, colloids):
      raise NotImplementedError("Implemented in child class.")

The above class contains two fundamental components, the kill switch and the call method.
As its name suggests, the kill switch is used to end the current run.
If this property is set to True,  will stop the current running engine and either restart it, or just finish.
This is used, for example, in episodic training, where the training should stop and the engine should reset after an amount of time or when a criterion is reached.
However, the kill switch can also be used for more critical cases, such as in an experiment when it appears that the agents would perform damaging actions.
The second component, the call function, is used to compute a value describing the current state of this task, ideally, large positive values mean the task is being achieved, whereas negative values mean it is actively not being achieved. 
For example, when training RL agents to rotate an object, the call method would compute the rotation speed and return some positive value if it was increasing or large, and a negative value if it were decreasing.

At times, we want agents to perform more than just one task, in these cases, two options exist.
The first is to write a custom task that provides feedback on two different objects such as how fast something is rotating and how far it has been pushed to a desired location.
However, these tasks often quickly become convoluted and difficult to reuse.
Therefore,  provides the  class which takes a list of single tasks as an argument and applies them to the engine.
This approach is more modular and allows for faster testing of strategies.



 §.§ Observables


A critical component of any RL algorithm and an important part of classical agent control approaches is their state description.
Rather than passing all state information (e.g., all positions and directions of the agents) to the decision making model, the observable usually selects from or condenses the information considerably.
This mimics the limited sensing capabilities that agents might have.
For example, agents might only be able to retrieve information about other particles within a finite sensing radius or might not be able to tell the direction of their neighbours.
Investigating the amount of information needed for an agent to still be able to perform its task is very relevant for the design of cost-efficient micro agents.
In , the description of the environment as sensed by the agents is computed by .
All observables implemented in  inherit from a simple parent class outlined in Code Sample <ref>.
[language=Python, label=lst:observable, caption=  parent class.]
class Observable:
  def compute_observable(self, colloids):
    raise NotImplementedError("Implemented in child class.")

The class requires implementing one method, , which takes a list of particle objects as input and returns the state description vector for each of them.
This state vector is then passed on to the models by the force function.
As with the  class, we often want to construct agents capable of sensing the world in multiple ways.
To achieve this, we also provide a  class which takes a list of observables and computes a single output.



 §.§ Intrinsic Reward

In RL agents can be rewarded intrinsically to either encourage the exploration behavior of an agent or to assist in building its knowledge about the environment <cit.>. 
As the name suggests, such a method is intrinsic to the agent and is therefore implemented as a separate sub-module  which the agent can be equipped with. 



  §.§.§ Random Network Distillation

Random network distillation (RND) was introduced in 2018 by <cit.> as an efficient method for exploring the state space of reinforcement learning agents.

It has since been studied more broadly for its application and data selection for ML applications <cit.> and its implications on understanding data requirements for neural networks <cit.>. 

 provides , a module implementing RND, defining an intrinsic reward for agents, encouraging them to explore unseen regions of their state space. 
Going beyond the original version from <cit.>,  provides two implementations of RND, differing in their capability to memorize previous states.
For both versions of RND, a default setup is provided in swarmrl/intrinsic_reward/rnd_configs.py. 
Utilizing such configuration the implementation of RND is shown in Code Sample <ref>. 
[label=lst:intrinsic_reward_rnd, caption=Random Network Distillation as Intrinsic Reward, language=Python]Name
rnd_config = RNDConfig(...)

ac_agent = srl.agents.ActorCriticAgent(
    ...,
    intrinsic_reward=RNDReward(rnd_config)
) 

The implementation of the RND backend is based on the ZnNL python package <cit.>, offering a flexible but simple interface for complex training algorithms for neural networks.



 §.§ Exploration Policy

In addition to the sampling strategy which can select actions from its distributions that do not have the largest probability,  also provides direct access to exploration policies to further guide the agents into new, unexplored regions of their environment.


  §.§.§ Random Exploration

The simplest form of exploration policy is the  class, which an exploration probability can parameterize, ζ and decay rate, ϵ.
The exploration rate dictates how likely it is for an agent to take an exploration action over a policy action, and the decay rate slowly decreases this probability as the training evolves by 

    ζ^' = exp^-ϵt/Tζ,

where t is simulation time, and T is the time for a single training episode.





















 §.§ Trainer


The trainer is module in  responsible for handling the updates of the models and their interaction with the environment while developing a policy.
In the simplest case, the trainer deploys the models in the environment, collects the rewards over an episode, and performs a network update, as outlined in Code Sample <ref>.
However, in many cases more is demanded during RL training and even how rewards are collected can change depending on the problem.
Therefore, in , we have implemented several variations of the .


  §.§.§ Continuous vs Episodic

A user must first choose whether to perform episodic or continuous training.
In episodic training, agents are deployed in the environment for a fixed number of steps or until some condition is met.
After this time, a reward is computed based on their final state or trajectory during the episode.
This reward is then used to update the networks before the environment is reset and the process is started again.
In this approach, the episodes should be long enough to allow the agents to achieve their tasks but short enough to maximize efficiency.
A similar but alternative approach is semi-episodic training, wherein the agents are updated continuously during the training but after some time, or again, once a condition is met, the environment is reset and the training continues.
These options are available through the  packaged into .
The , displayed in Code Samples <ref> and <ref>, allows users to define how often the environment should be reset as a function of RL-episodes, or to pass a task which can force an environment reset.

[label=lst:cont-trainers, caption=Continuous Training, language=Python]trainers
rl_trainer = srl.trainers.ContinuousTrainer(
    [protocol_1, protocol_2],
    loss,
)
rl_trainer.perform_rl_training(
    system_runner=srl.Engine(...),
    n_episodes=5000,
    episode_length=50,
)


[label=lst:ep-trainers, caption=Episodic Training, language=Python]trainers
rl_trainer = srl.trainers.EpisodicTrainer(
                [protocol_1, protocol_2],
                loss,
            )
rl_trainer.perform_rl_training(
    get_engine=engine_fn,
    n_episodes=10,
    reset_frequency=30,  # > 1 is semi-episodic
    episode_length=100,
)     


Continuous training is a paradigm in which agents are deployed into an environment and are updated after a fixed amount of time without resetting the environment.
This online training is useful for situations where training must occur in the real world or where the time it will take to achieve a task is not well known.
 offers the  class to perform this kind of training.
It should be noted that if a  is provided to the continuous trainer capable of resetting the environment, the training will end.
This can also be helpful in cases where an end to training can be strictly identified.













































§ PERFORMANCE


Reinforcement learning typically does not require expensive neural networks, and classical control algorithms are limited in complexity by the number of particles they control.
However, speed becomes critical when considering the number of iterations they have to go through during training and deployment.
Therefore,  has been built with a performance-first mentality, leveraging the JAX ecosystem <cit.> heavily.
JAX is a Python library that, among other things, re-implements the NumPy library <cit.> with multi-threading, GPU deployment, and XLA-enhanced JIT compile capabilities.
Most variables in  are written using JAX arrays or their PyTree counterparts.
PyTrees extend the functionality JAX has for NumPy objects to data structures like classes.
Therefore, most of the class objects including the  class can be placed onto a GPU or parallelized over.
Furthermore, large parts of code are JIT compiled to ensure optimal performance.



§ SOFTWARE PRACTICES


In order to ensure the longevity and continued usability of the  software, we adhere to several software practices.
 is thoroughly tested with unit, integration, and functional tests in place.
Overall, we have 87 % coverage on the code with all classes having at least one unit test.
The code follows the NumPy doc-string format and is extensively documented which is hosted at <https://swarmrl.github.io/SwarmRL.ai/>.
Before new code is added to the project, all tests must pass and new tests should be included.
The project is hosted on GitHub where all these standards are enforced automatically.
Furthermore, we require at least one review from a code-owner, that is, a member of the core  team, to accept new changes to the software.




§ VISUALIZATION


Visualization is an essential part of the scientific process, particularly in simulations where one is searching for emergent strategy.
It is not always clear from plots alone if the agents have learned what they should do.
For this reason,  interfaces closely with the ZnVis <cit.> visualization library.
ZnVis, by design, handles the output data from  and is developed alongside the project.
The package can use custom mesh files for the agents, export rendered videos, and capture still-frame renderings using the Mitsuba engine <cit.>.
Figure <ref> illustrates the visualization capabilities of the package used on  experiments.



§ CONCLUSION


We have introduced , a Python package aimed at accelerating research into control and understanding of microscopic active particles, whether biological or artificial.
 allows researchers to implement state-of-the art control algorithms driven by reinforcement learning or classical policies, in both simulated and experiment settings.
As a Python package,  has a familiar interface for many scientists, and is simple enough to use with limited background in software.
Beyond its functionality, we have built  to be performant, capable of running on many cores, GPUs, and scalable to HPC clusters.
Our current focus is extending  to handle more control algorithms, e.g., MPC, Q-learning, additional experiment implementations, and additional engines.
Overall, our goal is that  allow scientists to spend more time creating new ideas, building new technology, and changing the field, rather than implementing and staying up to date with state of-of-the art machine learning methods.


§ ACKNOWLEDGEMENTS

C.H and S.T acknowledge financial support from the German Funding Agency (Deutsche Forschungsgemeinschaft DFG) under Germany’s Excellence Strategy EXC 2075-390740016.
C.H and S.T acknowledge financial support from the German Funding Agency (Deutsche Forschungsgemeinschaft DFG) under the Priority Program SPP 2363, “Utilization and Development of Machine Learning for Molecular Applications - Molecular Machine Learning” Project No. 497249646.
C.H and C.L acknowledge funding by the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) under Project Number 327154368-SFB 1313.
The authors acknowledge support by the state of Baden-Württemberg through bwHPC and the German Research Foundation (DFG) through grant INST 35/1597-1 FUGG.
The authors acknowledge support from the Deutsche Forschungsgemeinschaft (DFG, German Research Foundation) Compute Cluster grant no. 492175459.
S.T would like to acknowledge Annalena Daniels for her detailed review and comments on the manuscript.




