SetupAgent and Environment Grounding
Explore the role of the SetupAgent in initializing environment contexts, loading configurations, and preparing outputs essential for a Eureka-like reward learning pipeline. Learn how this agent sets shared states and produces baseline rollouts to support downstream AI agents in adaptive system design.
We'll cover the following...
- Code walk-through
- Loading configuration files safely
- Defining the SetupAgent and its execution entry point
- Guarding against repeated setup
- Loading loop and task configuration
- Enforcing the configuration contract
- Initializing the run directory and leaderboard
- Creating the environment and extracting its source
- Persisting run metadata
- Initializing shared session state
- Generating the baseline rollout
- Signaling completion to ADK
- Summary
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 yamlfrom google.adk.agents import BaseAgentfrom google.adk.agents.invocation_context import InvocationContextfrom google.adk.events import Eventfrom typing import AsyncGeneratorfrom loguru import logger
Let’s unpack this step by step:
We import
BaseAgentbecause every executable unit in an ADK workflow must subclass it. This is what allows ADK to recognizeSetupAgentas a runnable agent.InvocationContextis how ADK passes execution state into the agent. This object gives us access to the shared session state that persists across agents.Eventis 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.AsyncGeneratorappears because ADK agents are asynchronous and stream events as they execute.yamlis used to load configuration files, andloggergives 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_sourcefrom ..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 ...