Fractal’s Data Science Agent: Arya

Autonomous AI agents have been a prominent area of research for several years. Researchers have been obsessing over creating these agents that can act autonomously, as they are considered a great step towards AGI (Artificial General Intelligence). So let’s first define these agents – “they are AI entities that can perceive changes in the environment, make decisions, and take actions, towards completing a given task”. Over the years, their evolution has undergone several stages, starting from simple rule-based agents to reinforcement learning (RL) based agents. We all know AlphaGo, an RL agent that could beat humans in the board game called Go. But these RL agents are difficult to train, and they usually lack generalization capability. With the recent evolution of Large Language Models (LLMs), AI agents are receiving great attention once again as these LLMs have demonstrated remarkable generalization capabilities and can potentially serve as a foundational framework for designing AI agents. Through the exemplar of our autonomous AI agent Arya, we will now delve into the world of AI agents and understand how they are able to mimic humans by sensing, adapting, and interacting in their assigned environments to achieve a goal.

Understanding agents – the general conceptual framework

Drawing parallels with human capabilities, AI agents can be conceptualized as comprised of three main modules that work within some environment [See Figure 1]. The first one is called the perception module, which is analogous to the human sensory system and enables sensing the environment (e.g., sensing the room temperature). The second module, analogous to the human brain, handles information processing, thinking, planning, and decision-making. It is integrated with memory and knowledge (e.g., making decisions if it needs to adjust the thermostat). The third module functions like human limbs and enables the execution of actions using tools within the environment, thereby changing the environment (e.g., updating the thermostat settings and, hence, the room temperature). An agent can autonomously achieve complex tasks by repeating the process of sensing the environment, making decisions, and acting.

Fractal’s data science agent: Arya Agent Framework

Our team has developed an internal standardized framework to streamline the creation of AI agents. The framework comprises of several modules [See Figure 2]. The first and foremost module is agent persona. The specialized agents have a persona that defines the role and capabilities of the agents (e.g., researcher, python coder). Depending on the role, the agent is allowed to take specific actions needed for that persona; hence, the toolkit is tailored for these specific actions or tasks. Since the agent needs to process information to make decisions, the foundational models, such as LLM, serve as the agent’s cognitive engine. On top of that, thinking strategy becomes essential. This could involve established strategies like the Chain-of-Thought or Tree-of-Thought or new strategies tailored to the agent’s persona.

Additionally, a memory component governs how the agent stores information, knowledge, logs, and feedback. Another important capability is to be able to interact with humans and take their feedback in real-time. Finally, an evaluation framework assesses how effectively and successfully the agent is functioning. These components are standardized to form a cohesive framework for agent building. One can leverage this framework to define and fill the modules based on the agent’s overall functionality objectives.

Introducing Arya

Arya is an autonomous AI agent specialized in solving data science problems and participating in Kaggle competitions in an end-to-end manner [See Figure 3]. Given a competition name, she navigates the Kaggle website, retrieves the competition details, understands the competition, proposes an ML model and approach, writes Python scripts, executes the scripts to train the model, obtains results on the test dataset, submits the results on Kaggle and then finally fetches the rank of the submission. Within the framework discussed previously, the Kaggle website and local development setup constitute the environment for Arya. Her brain generates a task list and decides which task to perform at a given stage. Tools (such as Kaggle APIs, custom functions, search APIs, etc.) enable Arya to perform actions to complete these tasks such as execution of scripts, debugging, modelling, fetch the ranks, etc. An action usually triggers a change in the environment that Arya perceives and decide on the next task accordingly. For example, if there is an error during code execution, Arya will observe that and perform the debugging step before going to the next step of submission. Arya can work in collaboration with human users, making it feasible for users to give feedback at some important stages during the process. For example, she waits for user confirmation on the proposed ML approach before writing the codes. Users can interact and guide her for any modification.

When comparing Arya with other agents, such as a market research agent or a software engineering agent, distinct personas are at the core. Although the standardized framework stays constant, the capabilities differ according to each agent’s persona. For instance, the academic research agent explores open-source repositories of publications and platforms, like Hugging Face, to provide users with a tailored learning experience, whereas Arya interacts with Kaggle and helps user to take part in the competitions. The framework’s adaptability ensures that tools match specific capabilities, maintaining flexibility.

Technical challenges and ethical considerations

During Arya’s development, numerous technical challenges shaped her evolution. One significant hurdle occurred early on when creating persona description. Notably, changes in input prompts led to significant output variations, highlighting the system’s sensitivity to prompt changes. The persona descriptions need to be balanced so that the agent doesn’t become too narrow in its functioning and at the same time, it shouldn’t go beyond the role. It needed extensive experimentation to develop the persona for Arya.

Another important obstacle was to address token limit issue, particularly in prolonged conversations or multiple rounds of competition submissions within a single session. This is a common issue with complex agents where there are numerous tasks to complete a goal. However, substantial progress has been achieved in overcoming this challenge, allowing for long sessions and multiple rounds of submissions in the same session.

Hallucination was yet another prevalent issue that is inherent to LLMs. Arya, at times, prematurely believed that the competition had concluded, unexpectedly halting ongoing tasks. Additionally, occasional looping behavior, where Arya repeated steps instead of progressing, posed a persistent challenge. We have now made a good progress in overcoming this problem through prompt engineering, human-in-the-loop capability, and appropriate guardrails.

On the other hand, Ethical AI concerns are spanned at multiple levels – individuals, companies, and society at large. It includes concerns for privacy, bias & fairness, security, quality and reliability, IP rights, job displacement. At the individual level, it involves ensuring the quality and reliability of generated code, focusing on addressing accountability for potential failures in exceptional situations. Another facet concerns the explainability of the code generated by the agent, emphasizing the need for transparency. Using LLMs introduces concerns about biases inherent in training data, impacting the generated code. Detecting and addressing these biases emerges as a significant issue.

Impact of agents on data science roles

In the near future, a notable acceleration in data science projects is anticipated with the involvement of agents, particularly in automating traditional machine learning approaches. While the growth trajectory seemed promising, certain data science roles demanding deeper thinking and domain knowledge were acknowledged to remain beyond immediate automation. Experts in specific fields are hard to replace with AI. Humans have a unique advantage in connecting ideas across different areas and solving problems with interdisciplinary perspectives.

Speaking about the impact of autonomous AI agents like Arya on roles such as data scientists does shift the dynamics. While these agents write a significant amount of code, it could transform the role of a data scientist into more of a quality assurance role — ensuring the correctness and reliability of the generated code. On a positive note, it could unburden data scientists from traditional, repetitive tasks and inspire them to focus on solving complex and novel problems.

Arya’s future in data science

The overarching goal is to develop a versatile data science agent capable of addressing a broader spectrum of complex data science problems in real life. The evolving role of Arya in data science started with a focus on Kaggle competitions. The initial stages aimed to assess Arya’s capabilities in tackling diverse ML challenges while competing with real data scientists in the real world. Subsequent stages involve enhancing her capabilities for high rankings in complex Kaggle competitions. The goal remained to have a versatile data science agent that can work in collaboration with humans in solving a broader range of data science problems beyond Kaggle competitions. The ongoing research initiatives span incorporating community feedback, learning from experience, customized LLM training, and refining thinking strategies to align with data scientist’s way of thinking.

Integration with the Fractal’s data science agent platform

Arya is expected to be technically integrated into Genesis ecosystem. Fractal’s Generative AI platform, offers foundational models (both open-source and proprietary) services accessible through APIs, applications tailored for specific tasks, and agents that may be hosted on the platform, utilizing foundational models like LLMs. The three interconnected services within the Fractal’s data science agent: Arya platform — models, applications, and agents — are designed to work synergistically, providing a comprehensive and cohesive generative AI experience.