Search⌘ K
AI Features

SetupAgent and Environment Grounding

Explore the implementation of SetupAgent within an AI reward learning pipeline. Understand how to load configurations, initialize environments, prepare output directories, and generate baseline rollouts. This lesson helps you grasp how to set up shared execution state critical for subsequent AI agent steps.

In this EUREKA-like pipeline, the SetupAgent is the first step in the execution flow. It initializes shared state that downstream agents depend on: loaded configuration, an artifacts directory, and the initial environment context. It also generates a baseline rollout to serve as a visual reference before reward evolution begins. In this lesson, we’ll walk through agents/setup_agent.py.

Code walk-through

We begin by importing everything the SetupAgent needs to run.

import yaml
from google.adk.agents import BaseAgent
from google.adk.agents.invocation_context import InvocationContext
from google.adk.events import Event
from typing import AsyncGenerator
from loguru import logger
Imports: Declaring what this agent is allowed to do

Let’s unpack this step by step:

  • We import BaseAgent because every executable unit in an ADK workflow must subclass it. This is what allows ADK to recognize SetupAgent as a runnable agent.

  • InvocationContext is how ADK passes execution state into the agent. This object gives us access to the shared session state that persists across agents.

  • Event is used to signal completion back to ADK. Even when an agent doesn’t emit meaningful output, it still needs to yield an event to satisfy the execution contract.

  • AsyncGenerator appears because ADK agents are asynchronous and stream events as they execute.

  • yaml is used to load configuration files, and logger gives us structured logging during setup.

At this point, we’ve declared:

“This file defines an ADK agent that reads configuration and mutates shared execution state.”

Next, we import the tools that the SetupAgent will rely on.

from ..tools.brax_env import get_env_and_source
from ..tools.artifacts import init_run_dir, init_leaderboard, save_json

These imports are especially important because they tell us what this agent does not do itself.

  • get_env_and_source(...) encapsulates ...