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: Un Marco de Modelado Basado en Agentes Basado en JAX

Información Básica

  • ID del Artículo: 2409.06345
  • Título: Foragax: An Agent-Based Modelling Framework Based on JAX
  • Autores: Siddharth Chaturvedi, Ahmed El-Gazzar, Marcel van Gerven (Universidad de Radboud)
  • Clasificación: cs.MA (Sistemas Multiagente)
  • Fecha de Publicación: Septiembre de 2024 (preimpresión en arXiv)
  • Enlace del Artículo: https://arxiv.org/abs/2409.06345
  • Repositorio de Código: https://github.com/i-m-iron-man/Foragax.git

Resumen

La búsqueda de alimento es una actividad universal realizada por organismos biológicos en entornos compartidos para mantener la homeostasis. Mediante el modelado computacional del comportamiento de búsqueda de alimento de múltiples agentes, es posible estudiar comportamientos emergentes individuales y colectivos de manera manejable. El modelado basado en agentes ha demostrado ser efectivo en la simulación de tales tareas, pero extender las simulaciones para acomodar grandes cantidades de agentes con dinámicas complejas sigue siendo desafiante. Este artículo presenta Foragax, un conjunto de herramientas universal, escalable y acelerado por hardware para modelado de búsqueda de alimento multiagente. Aprovechando la biblioteca JAX, el conjunto de herramientas puede simular el comportamiento de búsqueda de alimento de miles de agentes en un entorno común de manera vectorizada y diferenciable de extremo a extremo. El conjunto de herramientas proporciona herramientas de modelado de agentes para simular diversas tareas de búsqueda de alimento, incluyendo opciones para diseñar dinámicas de agentes espaciales y temporales personalizadas, estrategias de control, modelos de sensores y condiciones de límite. Además, el número de agentes puede aumentarse o disminuirse durante la simulación según reglas personalizadas. Aunque se aplica a la búsqueda de alimento, el conjunto de herramientas también puede utilizarse para modelar y simular una amplia gama de otros escenarios multiagente.

Antecedentes de Investigación y Motivación

Definición del Problema

La búsqueda de alimento es uno de los problemas de optimización de comportamiento más antiguos y críticos que todos los organismos biológicos deben resolver para sobrevivir. Desde las bacterias más simples hasta los primates altamente socializados, el instinto de una búsqueda de alimento eficiente es omnipresente en todas las etapas de la evolución. La búsqueda de alimento en la naturaleza es inherentemente una actividad multiagente, ya que los organismos adyacentes comparten el mismo entorno de búsqueda de alimento.

Limitaciones de los Enfoques Existentes

  1. Problemas de Escalabilidad: Los métodos tradicionales de modelado basado en agentes enfrentan cuellos de botella computacionales al tratar con grandes cantidades de agentes
  2. Falta de Realismo Biológico: Las limitaciones en el número de agentes en simulaciones existentes resultan en inconsistencias con los comportamientos de búsqueda de alimento observados en la naturaleza
  3. Complejidad Computacional: El estado de cada agente depende del comportamiento de todos los demás agentes, haciendo que la simulación sea extremadamente difícil a medida que aumenta el número de agentes

Motivación de la Investigación

Los autores desean proporcionar a los modeladores de tareas de búsqueda de alimento la capacidad de simular simultáneamente grandes cantidades (miles) de agentes en un único entorno. Esta capacidad de simulación a gran escala es crucial para comprender comportamientos colectivos emergentes complejos y para cerrar la brecha entre la investigación teórica y las observaciones naturales.

Contribuciones Principales

  1. Presentación del Marco Foragax: Un conjunto de herramientas universal, escalable y acelerado por hardware para modelado de búsqueda de alimento multiagente basado en JAX
  2. Implementación de Simulación Paralela a Gran Escala: Capacidad de simular miles de agentes mediante métodos de vectorización de extremo a extremo
  3. Provisión de Interfaz de Modelado Flexible: Soporte para dinámicas personalizadas espaciales y temporales, estrategias de control, modelos de sensores y condiciones de límite
  4. Soporte para Gestión Dinámica de Agentes: Capacidad de aumentar o disminuir el número de agentes durante la simulación según reglas personalizadas
  5. Implementación de Simulación Continua No Episódica: Soporte para simulaciones a largo plazo sin reiniciar el entorno y el estado de los agentes

Explicación Detallada del Método

Definición de la Tarea

Foragax se enfoca en simular el comportamiento de búsqueda de alimento de grandes cantidades de agentes en un entorno común mientras se mantiene la homeostasis. Las características de la tarea incluyen:

  • Entrada: Estado inicial de agentes, configuración del entorno, distribución de recursos
  • Salida: Trayectorias de comportamiento de agentes, patrones de comportamiento emergente colectivo
  • Restricciones: Limitación de recursos, interacciones entre agentes, condiciones de límite del entorno

Arquitectura del Marco

Foragax se basa en el diseño de clases con cinco niveles principales de abstracción:

1. Policy (Política)

  • Representada por un conjunto de variables de estado y parámetros
  • Utilizada para generar comportamiento de agentes en el entorno
  • Soporta múltiples modelos: modelo de Wilson-Cowan, redes neuronales recurrentes, perceptrones multicapa
  • Puede modelarse como determinista o estocástica

2. Agent (Agente)

  • Representado por un conjunto de variables de estado y parámetros
  • Soporta dinámicas temporales y espaciales personalizadas
  • Mantiene un índice único y estado activo (1 para activo, 0 para inactivo)
  • Las dinámicas pueden ser manipuladas por su política

3. Agent-set (Conjunto de Agentes)

  • Representación de conjunto de agentes
  • Todos los agentes comparten la misma estructura de variables de estado, parámetros y política
  • Soporta operaciones de selección, ordenamiento, adición o eliminación de agentes
  • Mantiene la forma de la estructura de datos constante mediante relleno con ceros, preservando las características de JIT y vectorización automática de JAX

4. Learning Algorithms (Algoritmos de Aprendizaje)

  • Proporciona un conjunto de algoritmos para ajustar parámetros de política y agentes
  • La implementación actual se basa en métodos de neuroevolución sin gradientes
  • Se planea expandir a métodos de aprendizaje por refuerzo basados en gradientes

5. Tasks (Tareas)

  • Permite diseñar entornos de búsqueda de alimento con condiciones de límite personalizadas y colocación de paredes
  • Utiliza algoritmos de trazado de rayos para detección de colisiones

Puntos de Innovación Técnica

1. Optimización del Backend JAX

  • Aprovecha la compilación justo a tiempo (JIT) y vectorización automática (AV) de JAX
  • Implementa cálculo vectorizado y diferenciable de extremo a extremo
  • Soporta aceleración por hardware (GPU/TPU)

2. Estrategia de Relleno con Ceros

  • Mediante preasignación de estructuras de datos de tamaño máximo
  • Utiliza relleno con ceros para mantener forma constante
  • Preserva características de optimización de JAX en operaciones de conjunto de agentes

3. Simulación Continua No Episódica

  • Sin reinicio del entorno y estado de agentes
  • Soporta fluctuaciones en el tamaño de la población
  • Facilita la emergencia de comportamientos adaptativos de final abierto

Configuración Experimental

Configuración de Experimentos

  • Número de Agentes: 1000 agentes
  • Escala de Red Neuronal: Modelo de población neuronal basado en tasas de 50 neuronas por agente
  • Número de Recursos: 300 parches de recursos
  • Pasos de Simulación: 1 millón de pasos
  • Plataforma Computacional: Computadora de clúster

Modelo de Dinámicas

Dinámicas de Recursos

El valor de los recursos sigue dinámicas de Lotka-Volterra:

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

Donde:

  • ϵ y α son las tasas de crecimiento y decaimiento de recursos respectivamente
  • w(xn, xm) representa una función inversamente proporcional a la distancia entre la ubicación del n-ésimo recurso xn y la ubicación del m-ésimo agente xm

Modelo de Movimiento de Agentes

La posición y velocidad de los agentes utilizan un modelo de integrador doble:

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

Donde:

  • qm = xm, ẋmᵀ es el vector de estado del agente
  • um ∈ ℝ es la entrada de aceleración
  • La aceleración se determina por el promedio de lecturas lineales de una red neuronal recurrente de 50 neuronas

Resultados Experimentales

Resultados Principales

  • Eficiencia Computacional: La simulación de 1000 agentes, 300 recursos y 1 millón de pasos se completa en 40 minutos
  • Verificación de Escalabilidad: Demostración exitosa del comportamiento de búsqueda de alimento de 600 agentes en 600 recursos
  • Estabilidad del Sistema: La distribución espacial de agentes y recursos se mantiene razonable durante la simulación

Desempeño

Los resultados experimentales demuestran que Foragax puede:

  1. Manejar eficientemente simulaciones multiagente a gran escala
  2. Mantener requisitos de rendimiento en tiempo real
  3. Generar patrones de comportamiento de búsqueda de alimento biológicamente razonables

Resultados de Visualización

El artículo proporciona visualizaciones de 600 agentes (puntos azules) buscando alimento dispersos entre 600 recursos (puntos rojos), con flechas que indican la dirección de los agentes, demostrando las características de organización espacial del sistema.

Trabajo Relacionado

Marcos de Modelado Basado en Agentes

  • Agents.jl: Software de modelado basado en agentes de alto rendimiento y rico en funcionalidades
  • Mesa: Marco de modelado basado en agentes en Python
  • Los marcos existentes presentan cuellos de botella de rendimiento en simulaciones a gran escala

Investigación en Teoría de Búsqueda de Alimento

  • Investigación del comportamiento de búsqueda de alimento desde bacterias simples hasta primates complejos
  • Desarrollo de teoría óptima de búsqueda de alimento
  • Investigación del impacto de reservas de energía internas en la capacidad de respuesta a recompensas

Aplicaciones de Sistemas Multiagente

  • Modelado de propagación de enfermedades
  • Simulación de comportamiento de enjambres
  • Problemas de optimización de programación

Conclusiones y Discusión

Conclusiones Principales

  1. Foragax implementa exitosamente simulación eficiente de comportamiento de búsqueda de alimento multiagente a gran escala
  2. El backend JAX mejora significativamente el rendimiento computacional y la escalabilidad
  3. El marco proporciona suficiente flexibilidad para modelar diversos escenarios de búsqueda de alimento
  4. La simulación continua no episódica puede revelar comportamientos emergentes interesantes

Limitaciones

  1. Limitaciones de Preasignación: Requiere conocimiento previo del tamaño máximo de las estructuras de datos
  2. Consumo de Memoria: La estrategia de relleno con ceros puede resultar en uso ineficiente de memoria
  3. Alcance de Algoritmos: Actualmente soporta principalmente algoritmos de neuroevolución, con soporte de aprendizaje por refuerzo pendiente
  4. Validación Insuficiente: Falta de comparación cuantitativa con comportamiento real de búsqueda de alimento biológico

Direcciones Futuras

  1. Expansión a algoritmos de aprendizaje por refuerzo basados en gradientes
  2. Mejora de la eficiencia de uso de memoria
  3. Aumento de validación de realismo biológico
  4. Desarrollo de herramientas de análisis más ricas

Evaluación Profunda

Fortalezas

  1. Fuerte Innovación Técnica: Introduce las capacidades de computación de alto rendimiento de JAX al campo del modelado basado en agentes
  2. Alto Valor Práctico: Resuelve un cuello de botella clave en simulación de agentes a gran escala
  3. Diseño Razonable: Arquitectura de abstracción de cinco niveles clara, fácil de extender y usar
  4. Rendimiento Excelente: Logra mejoras de rendimiento significativas mientras mantiene funcionalidad

Deficiencias

  1. Validación Experimental Limitada: Principalmente demuestra viabilidad técnica, carece de análisis de comportamiento profundo
  2. Validación Biológica Insuficiente: Sin comparación cuantitativa con comportamiento real de búsqueda de alimento
  3. Contribución Teórica Limitada: Principalmente implementación de ingeniería, con innovación teórica relativamente limitada
  4. Documentación Incompleta: Ciertos detalles técnicos carecen de descripción suficiente

Impacto

  1. Contribución Disciplinaria: Proporciona una herramienta computacional poderosa para investigación de sistemas multiagente
  2. Valor Práctico: Puede soportar investigación ecológica y conductual a mayor escala y más compleja
  3. Reproducibilidad: El código de código abierto mejora la reproducibilidad y extensibilidad de la investigación
  4. Aplicación Interdisciplinaria: Aunque enfocado en búsqueda de alimento, es extensible a otros escenarios multiagente

Escenarios Aplicables

  1. Investigación Ecológica: Modelado de comportamiento de búsqueda de alimento de animales a gran escala
  2. Biología Evolutiva: Investigación de dinámicas evolutivas de poblaciones
  3. Inteligencia Artificial: Prueba de algoritmos de aprendizaje por refuerzo multiagente
  4. Sistemas Complejos: Investigación de comportamientos emergentes y fenómenos de autoorganización
  5. Ciencias Sociales: Modelado de dinámicas sociales a gran escala

Referencias

El artículo cita literatura clave de múltiples disciplinas importantes, incluyendo:

  • Documentación técnica de la biblioteca JAX
  • Obras clásicas en teoría de búsqueda de alimento (Stephens & Krebs, 1986)
  • Marcos de modelado basado en agentes existentes (Agents.jl, Mesa)
  • Avances recientes en neuroevolución y aprendizaje por refuerzo

Evaluación General: Este es un artículo de alta calidad que, mediante la combinación de tecnología de computación de alto rendimiento moderna con modelado basado en agentes, resuelve un cuello de botella técnico importante en el campo. Aunque relativamente limitado en innovación teórica, su valor de ingeniería y practicidad lo convierten en una contribución importante a la investigación de sistemas multiagente.