2025-11-22T09:37:22.788326

Foragax: An Agent-Based Modelling Framework Based on JAX

Chaturvedi, El-Gazzar, van Gerven
Foraging for resources is a ubiquitous activity conducted by living organisms in a shared environment to maintain their homeostasis. Modelling multi-agent foraging in-silico allows us to study both individual and collective emergent behaviour in a tractable manner. Agent-based modelling has proven to be effective in simulating such tasks, though scaling the simulations to accommodate large numbers of agents with complex dynamics remains challenging. In this work, we present Foragax, a general-purpose, scalable, hardware-accelerated, multi-agent foraging toolkit. Leveraging the JAX library, our toolkit can simulate thousands of agents foraging in a common environment, in an end-to-end vectorized and differentiable manner. The toolkit provides agent-based modelling tools to model various foraging tasks, including options to design custom spatial and temporal agent dynamics, control policies, sensor models, and boundary conditions. Further, the number of agents during such simulations can be increased or decreased based on custom rules. While applied to foraging, the toolkit can also be used to model and simulate a wide range of other multi-agent scenarios.
academic

Foragax: Фреймворк агентного моделирования на основе JAX

Основная информация

  • ID статьи: 2409.06345
  • Название: Foragax: An Agent-Based Modelling Framework Based on JAX
  • Авторы: Siddharth Chaturvedi, Ahmed El-Gazzar, Marcel van Gerven (Radboud University)
  • Классификация: cs.MA (Многоагентные системы)
  • Дата публикации: Сентябрь 2024 г. (препринт arXiv)
  • Ссылка на статью: https://arxiv.org/abs/2409.06345
  • Репозиторий кода: https://github.com/i-m-iron-man/Foragax.git

Аннотация

Поиск пищи является универсальной деятельностью организмов для поддержания гомеостаза в общей среде. Посредством вычислительного моделирования многоагентного поведения при поиске пищи можно исследовать возникающее поведение на индивидуальном и коллективном уровнях управляемым образом. Агентное моделирование доказало свою эффективность при моделировании таких задач, однако расширение моделирования для размещения большого количества агентов со сложной динамикой остаётся сложной задачей. В данной статье представляется Foragax — универсальный, масштабируемый, аппаратно-ускоренный набор инструментов для многоагентного поиска пищи. Используя библиотеку JAX, набор инструментов может моделировать поведение при поиске пищи тысяч агентов в общей среде полностью векторизованным и дифференцируемым способом. Набор инструментов предоставляет средства агентного моделирования для моделирования различных задач поиска пищи, включая возможности проектирования пользовательской пространственной и временной динамики агентов, стратегий управления, моделей датчиков и граничных условий. Кроме того, количество агентов может быть увеличено или уменьшено во время моделирования в соответствии с пользовательскими правилами. Хотя применяется к поиску пищи, набор инструментов также может использоваться для моделирования и симуляции широкого спектра других многоагентных сценариев.

Исследовательский контекст и мотивация

Определение проблемы

Поиск пищи является одной из самых древних и критических задач оптимизации поведения, которую должны решать все организмы для выживания. От простейших бактерий до высокосоциализированных приматов, инстинкт эффективного поиска пищи распространён на всех этапах эволюции. Поиск пищи в природе по своей сути является многоагентной деятельностью, поскольку соседние организмы разделяют одну и ту же среду поиска пищи.

Ограничения существующих подходов

  1. Проблемы масштабируемости: Традиционные методы агентного моделирования сталкиваются с вычислительными узкими местами при работе с большим количеством агентов
  2. Недостаточная биологическая реалистичность: Ограничение количества агентов в существующих моделях приводит к несоответствиям с поведением при поиске пищи, наблюдаемым в природе
  3. Вычислительная сложность: Состояние каждого агента зависит от поведения всех остальных агентов, что делает моделирование чрезвычайно сложным с увеличением количества агентов

Исследовательская мотивация

Авторы стремились предоставить разработчикам моделей поиска пищи возможность одновременного моделирования большого количества (тысяч) агентов в единой среде. Такая крупномасштабная способность моделирования критически важна для понимания сложного возникающего коллективного поведения и сокращения разрыва между теоретическими исследованиями и естественными наблюдениями.

Основные вклады

  1. Предложен фреймворк Foragax: Универсальный, масштабируемый, аппаратно-ускоренный набор инструментов для многоагентного моделирования поиска пищи на основе JAX
  2. Реализовано крупномасштабное параллельное моделирование: Способность моделировать тысячи агентов посредством полностью векторизованного подхода
  3. Предоставлен гибкий интерфейс моделирования: Поддержка пользовательской пространственной и временной динамики, стратегий управления, моделей датчиков и граничных условий
  4. Поддержка динамического управления агентами: Возможность увеличения или уменьшения количества агентов во время моделирования в соответствии с пользовательскими правилами
  5. Реализовано непрерывное неэпизодическое моделирование: Поддержка долгосрочного моделирования без сброса состояния среды и агентов

Подробное описание методов

Определение задачи

Foragax сосредоточен на моделировании поведения при поиске пищи большого количества агентов в общей среде при одновременном поддержании гомеостаза. Характеристики задачи включают:

  • Входные данные: Начальное состояние агентов, конфигурация среды, распределение ресурсов
  • Выходные данные: Траектории поведения агентов, модели возникающего коллективного поведения
  • Ограничения: Ограниченность ресурсов, взаимодействие между агентами, граничные условия среды

Архитектура фреймворка

Foragax основан на проектировании классов с пятью основными уровнями абстракции:

1. Policy (Стратегия)

  • Представлена набором переменных состояния и параметров
  • Используется для создания поведения агентов в среде
  • Поддерживает несколько моделей: модель Вильсона-Коуэна, рекуррентные нейронные сети, многослойные персептроны
  • Может моделироваться как детерминированная или стохастическая

2. Agent (Агент)

  • Представлен набором переменных состояния и параметров
  • Поддерживает пользовательскую временную и пространственную динамику
  • Сохраняет уникальный индекс и статус активности (1 для активного, 0 для неактивного)
  • Динамика может управляться его стратегией

3. Agent-set (Множество агентов)

  • Коллективное представление агентов
  • Все агенты разделяют одну и ту же структуру переменных состояния, параметров и стратегии
  • Поддерживает операции выбора, сортировки, добавления или удаления агентов
  • Посредством нулевого заполнения сохраняет постоянную форму структуры данных, поддерживая JIT и автоматическую векторизацию JAX

4. Learning Algorithms (Алгоритмы обучения)

  • Предоставляет набор алгоритмов для настройки параметров стратегии и агентов
  • Текущая реализация основана на методах нейроэволюции без градиентов
  • Планируется расширение на методы обучения с подкреплением на основе градиентов

5. Tasks (Задачи)

  • Позволяет проектировать среды поиска пищи с пользовательскими граничными условиями и расположением стен
  • Использует алгоритм трассировки лучей для обнаружения столкновений

Технические инновации

1. Оптимизация JAX-бэкенда

  • Использование немедленной компиляции (JIT) и автоматической векторизации (AV) JAX
  • Реализация полностью векторизованных и дифференцируемых вычислений
  • Поддержка аппаратного ускорения (GPU/TPU)

2. Стратегия нулевого заполнения

  • Предварительное выделение структур данных максимального размера
  • Использование нулевого заполнения для сохранения постоянной формы
  • Сохранение оптимизационных свойств JAX при операциях с множествами агентов

3. Непрерывное неэпизодическое моделирование

  • Отсутствие сброса состояния среды и агентов
  • Поддержка колебаний размера популяции
  • Содействие возникновению открытого адаптивного поведения

Экспериментальная установка

Конфигурация экспериментов

  • Количество агентов: 1000 агентов
  • Размер нейронной сети: Модель нейронной популяции на основе скорости с 50 нейронами на агента
  • Количество ресурсов: 300 ресурсных участков
  • Количество шагов моделирования: 1 миллион шагов
  • Вычислительная платформа: Кластерный компьютер

Модели динамики

Динамика ресурсов

Стоимость ресурса следует динамике Лотки-Вольтерра:

ṡn = ϵsn - αs²n - Σ w(xn, xm)

где:

  • ϵ и α — коэффициенты роста и затухания ресурса соответственно
  • w(xn, xm) — функция, обратно пропорциональная расстоянию между n-м местоположением ресурса xn и m-м местоположением агента xm

Модель движения агента

Положение и скорость агента используют модель двойного интегратора:

q̇m = [0 1; 0 0]qm + [0; 1]um

где:

  • qm = xm, ẋmᵀ — вектор состояния агента
  • um ∈ ℝ — входное ускорение
  • Ускорение определяется средним значением линейного считывания из рекуррентной нейронной сети из 50 нейронов

Результаты экспериментов

Основные результаты

  • Вычислительная эффективность: Моделирование 1000 агентов, 300 ресурсов и 1 миллиона шагов завершается за 40 минут
  • Проверка масштабируемости: Успешная демонстрация поведения при поиске пищи 600 агентов среди 600 ресурсов
  • Стабильность системы: Пространственное распределение агентов и ресурсов остаётся разумным на протяжении моделирования

Производительность

Результаты экспериментов показывают, что Foragax способен:

  1. Эффективно обрабатывать крупномасштабное многоагентное моделирование
  2. Поддерживать требования к производительности в реальном времени
  3. Генерировать биологически обоснованные модели поведения при поиске пищи

Результаты визуализации

Статья предоставляет визуализацию рассеянного поиска пищи 600 агентами (синие точки) среди 600 ресурсов (красные точки), где стрелки указывают направление агентов, демонстрируя пространственные организационные характеристики системы.

Связанные работы

Фреймворки агентного моделирования

  • Agents.jl: Высокопроизводительное, функционально богатое программное обеспечение для агентного моделирования
  • Mesa: Фреймворк агентного моделирования на Python
  • Существующие фреймворки имеют узкие места производительности при крупномасштабном моделировании

Исследования теории поиска пищи

  • Исследования поведения при поиске пищи от простых бактерий до сложных приматов
  • Развитие теории оптимального поиска пищи
  • Исследования влияния внутренних энергетических запасов на чувствительность к вознаграждению

Приложения многоагентных систем

  • Моделирование распространения заболеваний
  • Моделирование группового поведения
  • Задачи оптимизации расписания

Выводы и обсуждение

Основные выводы

  1. Foragax успешно реализует эффективное моделирование крупномасштабного многоагентного поведения при поиске пищи
  2. JAX-бэкенд значительно повышает вычислительную производительность и масштабируемость
  3. Фреймворк обеспечивает достаточную гибкость для моделирования различных сценариев поиска пищи
  4. Непрерывное неэпизодическое моделирование способно выявить интересное возникающее поведение

Ограничения

  1. Ограничения предварительного выделения: Необходимо заранее знать максимальный размер структур данных
  2. Потребление памяти: Стратегия нулевого заполнения может привести к неэффективному использованию памяти
  3. Диапазон алгоритмов: Текущая поддержка в основном охватывает алгоритмы нейроэволюции, поддержка обучения с подкреплением требует совершенствования
  4. Недостаточная валидация: Отсутствие количественного сравнения с реальным поведением животных при поиске пищи

Будущие направления

  1. Расширение на алгоритмы обучения с подкреплением на основе градиентов
  2. Улучшение эффективности использования памяти
  3. Увеличение валидации биологической реалистичности
  4. Разработка более богатого набора инструментов анализа

Глубокая оценка

Преимущества

  1. Сильная техническая инновативность: Внедрение высокопроизводительных вычислительных возможностей JAX в область многоагентного моделирования
  2. Высокая практическая ценность: Решение критических узких мест в крупномасштабном моделировании агентов
  3. Рациональное проектирование: Ясная пятиуровневая архитектура абстракции, легко расширяемая и используемая
  4. Отличная производительность: Достижение значительного повышения производительности при сохранении функциональности

Недостатки

  1. Ограниченная экспериментальная валидация: Демонстрация в основном технической осуществимости, отсутствие глубокого анализа поведения
  2. Недостаточная биологическая валидация: Отсутствие количественного сравнения с реальным поведением при поиске пищи
  3. Ограниченный теоретический вклад: В основном инженерная реализация, относительно небольшие теоретические инновации
  4. Недостаточно подробная документация: Некоторые технические детали описаны недостаточно полно

Влияние

  1. Вклад в область: Предоставление мощного вычислительного инструмента для исследований многоагентных систем
  2. Практическая ценность: Возможность поддержки крупномасштабных и более сложных исследований в экологии и поведении
  3. Воспроизводимость: Открытый исходный код повышает воспроизводимость и расширяемость исследований
  4. Междисциплинарные приложения: Хотя сосредоточено на поиске пищи, может быть расширено на другие многоагентные сценарии

Применимые сценарии

  1. Экологические исследования: Моделирование крупномасштабного поведения животных при поиске пищи
  2. Эволюционная биология: Исследование динамики популяционной эволюции
  3. Искусственный интеллект: Тестирование алгоритмов многоагентного обучения с подкреплением
  4. Сложные системы: Исследование возникающего поведения и явлений самоорганизации
  5. Социальные науки: Моделирование крупномасштабной социальной динамики

Библиография

Статья цитирует ключевые работы из нескольких важных областей, включая:

  • Техническую документацию библиотеки JAX
  • Классические труды по теории поиска пищи (Stephens & Krebs, 1986)
  • Существующие фреймворки агентного моделирования (Agents.jl, Mesa)
  • Последние достижения в нейроэволюции и обучении с подкреплением

Общая оценка: Это высококачественная системная статья, которая решает важное техническое узкое место в этой области путём объединения современных технологий высокопроизводительных вычислений с многоагентным моделированием. Хотя относительно ограничена в теоретических инновациях, её инженерная ценность и практическая полезность делают её значительным вкладом в исследования многоагентных систем.