DOCSEARCHbyAlgolia
Submit Docs

Aviary Gymnasium

repository·main·Indexed 43 days ago

https://github.com/future-house/aviary

A gymnasium for defining custom language agent Reinforcement Learning (RL) environments, designed to work with the Language Decision Processes (LDP) framework. Aviary provides pre-existing environments for math (GSM8k), general knowledge (HotPotQA), scientific literature search (LAB-Bench, LFRQA), and Jupyter notebook execution.

Tokens
26.1K
Snippets
80
Records
105
Agent score
54%

What's inside Aviary

  1. Overview of Aviary

    main

    Aviary is a gymnasium designed for defining custom language agent Reinforcement Learning (RL) environments. It is intended to work alongside the LDP (Language Decision Processes) library, which allows users to define custom language agents.

    Aviary provides several pre-existing environments across different domains:

    • Math
    • General knowledge
    • Biological sequences
    • Scientific literature search
    • Protein stability
    • Jupyter notebook execution (via packages/notebook)

  3. Implement a custom Environment

    main

    An Aviary environment must implement two primary asynchronous methods:

    1. await env.reset(): Initializes the environment, returns a list of initial observation messages (obs_msgs) and a list of available tools.
    2. await env.step(action_msg): Takes an action_msg (an instance of ToolRequestMessage) and returns (new_obs_msgs, reward, done, truncated).

    Communication is handled via Message objects. The action_msg contains tool calls, while obs_msgs contain general observations or ToolResponseMessage instances.

    obs_msgs, tools = await env.reset()
new_obs_msgs, reward, done, truncated = await env.step(action_msg)

  4. Configure containerized sandboxes for notebooks

    main

    For the recommended setup using containerized sandboxes, you must build the default Docker image and configure the environment variables to enable Docker usage.

    1. Build the image: Navigate to the docker/ directory and build the pinned image.

      cd docker/
docker build -t aviary-notebook-env -f Dockerfile.pinned .

    2. Set environment variables:

      • To enable Docker: Set NB_ENVIRONMENT_USE_DOCKER=true.
      • To use a custom image: If you are not using the default aviary-notebook-env image, override NB_ENVIRONMENT_DOCKER_IMAGE with your image name.
    # Build the default image
cd docker/
docker build -t aviary-notebook-env -f Dockerfile.pinned .

# Set environment variables (example for shell)
export NB_ENVIRONMENT_USE_DOCKER=true
# export NB_ENVIRONMENT_DOCKER_IMAGE=your-custom-image-name

  6. Create Functional Environments

    main

    For a lighter-weight approach, use the fenv functional interface. This uses decorators to define environments without explicit class subclassing.

    1. Use @fenv.start() to define the environment entry point. The function should return the initial observation and an initial state dictionary.
    2. Use @my_env.tool() to register tools. These tools can access the state dictionary passed in the start() function.
    3. Terminate the episode by setting state.done = True or setting a reward via state.reward.
    from aviary.core import fenv

@fenv.start()
def my_env(topic):
    # Returns (initial_observation, initial_state)
    return f"Write a story about {topic}", {}

@my_env.tool()
def multiply(x: float, y: float) -> float:
    """Multiply two numbers."""
    return x * y

@my_env.tool()
def print_story(story: str | bytes, state) -> None:
    """Print a story to the user and complete episode."""
    print(story)
    state.reward = 1
    state.done = True

# Usage
env = my_env(topic="foo")
obs, tools = await env.reset()

  8. Defining a Custom Environment in Aviary

    main

    To create a custom environment, inherit from aviary.core.Environment and specify your state type using a generic. You must implement reset() and step() methods.

    Key Methods:

    • reset(): Initializes the environment state. It should return a tuple containing a list of Message objects (the initial observation) and the list of available Tool objects.
    • step(action: ToolRequestMessage): Processes a tool call requested by the agent. It should return a tuple of (observations, reward, terminated, truncated).

    Core Components:

    • Environment[State]: The base class for all environments.
    • Message: Used to communicate content to the agent.
    • Tool: Created from functions using Tool.from_function(func) to allow agents to interact with the environment.
    • ToolRequestMessage: The type of action passed to the step method.
    from collections import namedtuple
from aviary.core import Environment, Message, ToolRequestMessage, Tool

# Define your state structure
CounterEnvState = namedtuple("CounterEnvState", ["count"])

class CounterEnv(Environment[CounterEnvState]):
    async def reset(self):
        # Initialize state
        self.state = CounterEnvState(count=0)
        self.target = 10
        
        # Define tools
        self.tools = [
            Tool.from_function(self.incr),
            Tool.from_function(self.decr),
        ]

        # Return initial observation and tools
        return [Message(content=f"Count to 10. counter={self.state.count}")], self.tools

    async def step(self, action: ToolRequestMessage):
        # Execute tool calls
        obs = await self.exec_tool_calls(action)

        # Calculate reward
        reward = int(self.state.count == self.target)

        # Return (observations, reward, terminated, truncated)
        return obs, reward, reward == 1, False

    def incr(self):
        self.state.count += 1
        return f"counter={self.state.count}"

    def decr(self):
        self.state.count -= 1
        return f"counter={self.state.count}"

  12. Install Aviary

    main

    You can install the core fhaviary package via pip. Depending on your needs, you can install just the core, the core with specific pre-existing environments, or the development version which includes tutorial notebooks.

    Core installation:

    pip install fhaviary

    Installation with environments: To include specific environments like gsm8k, hotpotqa, labbench, lfrqa, or notebook, use the following syntax:

    pip install 'fhaviary[gsm8k,hotpotqa,labbench,lfrqa,notebook]'

    Development installation: To install the package along with the tutorial notebooks:

    pip install "fhaviary[dev]"
    pip install fhaviary