AG2 is providing the key tools required to move towards automated AI-driven Cosmology and Astrophysics research. Using AG2, we have built a system that can autonomously carry out the key tasks involved in research workflows. Our fully open-source research-ready work, based on a Planning and Control strategy, has been successfully applied on real-world cosmology tasks.
— Boris Bolliet, Agentic AI Lead, Infosys-Cambridge AI Centre
Watch CMBAGENT in Action
Overview
Our AG2 application is called CMBAGENT. It's a multi-agent system for science with foundation models backends (e.g., Large Language Models; LLMs). It is built with the goal of having an AI system for Autonomous Scientific Discovery, with a focus on our research in Cosmology, which is the domain of origin of our team.
In Cosmology, our goal is to process data from telescopes to learn about the fundamental properties of the universe. This is an example of an "inverse problem" known as parameter inference, where we are given the data and, based on assumptions on how the universe works (expansion of space, formation of galaxies), we extract the most likely values of our model parameters.
The AG2 framework has allowed CMBAGENT to automate the workflow needed to solve the cosmological parameter inference problem.
Challenge
Measuring the fundamental properties of the universe based on telescope data requires running computationally expensive simulations and comparing their output with observations. The software packages required to run the simulations and confront them with measurements are research-level libraries that take PhD students several years to learn.
To automate these research challenges we are designing AI agents that can act as expert users on the libraries of interest and that can inform the analysis at hand with results from the scientific literature.
Solution: AG2 Integration
The very first version of CMBAGENT allowed us to exactly reproduce a cutting-edge cosmological data analysis in about 100 times less time and without writing the data analysis pipeline ourselves. The entire codebase was written by CMBAGENT.
We quickly realized that most of human feedback could in fact be easily substituted by adding more agents to the system, such as reviewer agents that critique the outputs and iterate, as well as adding more robustness with structured output and context awareness.
Planning and Control Strategy
In this strategy, a plan is first designed through a conversation between a planner agent and a plan reviewer agent. Once the user-specified number of rounds of reviews is exhausted, the plan is sent for execution to a control agent. During the Control phase, the control agent assigns the plan sub-tasks to the relevant agents for execution until all steps are successfully completed.
Key AG2 Capabilities Used
-
Multi-LLM calls: We utilize different LLM for different tasks. Current examples include gpt-4o-mini for simple tasks, claude-3.7 for reviewing, o3-mini-high for reasoning, gpt-4o for planning.
-
Swarm orchestration: To handle agent transitions we implemented CMBAGENT using swarm orchestration, where transitions are done based on conditions on a set of context variables that evolve and define the state of each agent.
-
Tool calls: Allow us to operate on data, either by executing code locally or storing relevant information into context.
-
Structured Output: A key feature that dramatically increases the robustness of CMBAGENT. During the planning stage, the plan has a very specific structure with steps corresponding to sub-tasks.
-
Retrieval Augmented Generation agents: The GPTAssistantAgent in AG2 is a simple and powerful interface to OpenAI assistant with vector stores containing research software packages, documentations, tutorials and scientific papers.
What's Next
CMBAGENT is only the start of our research programme and we consider it as a prototype for future systems that will be able to carry out scientific research with super-human capabilities.
-
Evaluation with Inspect.AI: We are developing an evaluation framework consisting of extensive benchmarks for cosmology research.
-
End-to-end cosmology research with ASTROPILOT: We have integrated CMBAGENT within ASTROPILOT, implementing not only the analysis, but also research idea generation and paper writing.
Conclusion
We think that in the not-so-distant future most fields in fundamental research will operate with self-driving laboratories. In Cosmology, this will look like teams of AI scientists collaborating to extract useful information from state-of-the-art datasets from telescopes.
Multi-agent frameworks such as AG2 will contribute to realise this paradigm shift in scientific research.
