Building a Eureka-Style Reward Loop with ADK

import gymnasium as gym
import numpy as np

class BalanceEnv(gym.Env):
    """Toy: balance a pole on a cart; minimise pole angle and cart drift."""

    # State variable spec the LLM gets to read.
    STATE_SPEC = {
        "cart_position": "float in [-2.4, 2.4]; episode ends if outside",
        "cart_velocity": "float; horizontal velocity",
        "pole_angle": "float in radians; vertical is 0.0",
        "pole_angular_velocity": "float; positive = falling right",
    }

    # The task-success metric — separate from the reward; LLM doesn't see this.
    TASK_SUCCESS = "fraction of episode steps where abs(pole_angle) < 0.1"

    # ... standard Gymnasium env body ...

REWARD_GEN_TEMPLATE = """You are designing a reward function for an RL agent.

TASK: {task}

The environment exposes these state variables (names you may reference):
{state_spec}

Write a Python function with this exact signature:

    def reward(state: dict[str, float], action: float, info: dict) -> tuple[float, dict[str, float]]:
        ...
        return total_reward, component_breakdown

`component_breakdown` is a dict from reward-component name to its
contribution that step. You must include at least 2 named components
and at most 5; the total is their weighted sum.

Constraints:
- Pure function. No globals, no imports beyond `math`.
- All branches must produce a finite float.
- Do not reference state variables outside the list above.

{prior_section}

Return the function as a single Python code block. No prose."""

PRIOR_SECTION_TEMPLATE = """The previous generation's best candidate is below.
Its reward components had these per-step statistics over training:

{stats}

The reflection on this candidate was:
{critique}

Improve on it. The new candidate must be different in structure
(not a parameter tweak) and address the reflection."""

import asyncio
from google.adk.agents import LlmAgent
from pathlib import Path

GEN_AGENT = LlmAgent(
    name="reward-generator",
    model="gemini-2.5-pro",
    instruction="Generate a reward function. Output Python only.",
)

async def generate_candidates(prompt: str, n: int) -> list[str]:
    """Generate n candidate reward functions in parallel."""
    async def one() -> str:
        result = await GEN_AGENT.arun(prompt)
        return _extract_code_block(result.final_response)

    return await asyncio.gather(*(one() for _ in range(n)))


def _extract_code_block(text: str) -> str:
    """Pull the python code block out of the LLM response."""
    import re
    m = re.search(r"```(?:python)?\s*\n(.*?)```", text, re.DOTALL)
    if not m:
        raise ValueError(f"no code block in response: {text[:200]}")
    return m.group(1).strip()

import ast
import math

ALLOWED_NAMES = {"math", "abs", "min", "max", "float", "int"}

def validate_reward_source(source: str) -> tuple[bool, str]:
    """Static checks before we run the candidate."""
    try:
        tree = ast.parse(source)
    except SyntaxError as e:
        return False, f"syntax: {e}"

    if not any(isinstance(n, ast.FunctionDef) and n.name == "reward" for n in ast.walk(tree)):
        return False, "missing function named 'reward'"

    # Reject imports beyond math
    for node in ast.walk(tree):
        if isinstance(node, ast.Import) and any(a.name != "math" for a in node.names):
            return False, "disallowed import"
        if isinstance(node, ast.ImportFrom) and node.module != "math":
            return False, "disallowed from-import"

    return True, ""


def compile_reward(source: str):
    """Compile and return the reward function or raise."""
    ok, msg = validate_reward_source(source)
    if not ok:
        raise ValueError(f"invalid reward source: {msg}")

    namespace = {"math": math}
    exec(source, namespace)
    fn = namespace.get("reward")
    if not callable(fn):
        raise ValueError("no callable named 'reward' after exec")
    return fn

import gymnasium as gym
import numpy as np
from stable_baselines3 import PPO

def wrap_env_with_reward(base_env, reward_fn):
    """Wrap a Gymnasium env so its reward comes from `reward_fn`."""
    class _Wrap(gym.Wrapper):
        def step(self, action):
            obs, _, terminated, truncated, info = self.env.step(action)
            state = self._obs_to_state(obs)
            r, components = reward_fn(state, float(action), info)
            info["reward_components"] = components
            return obs, r, terminated, truncated, info

        def _obs_to_state(self, obs):
            return {
                "cart_position": float(obs[0]),
                "cart_velocity": float(obs[1]),
                "pole_angle": float(obs[2]),
                "pole_angular_velocity": float(obs[3]),
            }

    return _Wrap(base_env)


def train_and_eval(reward_source: str, total_steps: int = 50_000):
    """Train a policy with this reward; return (success_score, component_stats)."""
    reward_fn = compile_reward(reward_source)
    env = wrap_env_with_reward(gym.make("CartPole-v1"), reward_fn)

    component_log: dict[str, list[float]] = {}

    def _log_components(info):
        for k, v in info.get("reward_components", {}).items():
            component_log.setdefault(k, []).append(float(v))

    model = PPO("MlpPolicy", env, verbose=0)
    model.learn(total_timesteps=total_steps, callback=lambda *a, **k: None)

    # Evaluate against the held-out task-success metric
    success = _eval_policy(model, env)
    stats = _summarise(component_log)
    return success, stats

import asyncio, json
from pathlib import Path

async def run_eureka_loop(
    task: str,
    state_spec: dict,
    generations: int = 5,
    candidates_per_gen: int = 8,
    out_dir: Path = Path("runs"),
):
    out_dir.mkdir(exist_ok=True)
    prior_critique = None
    prior_best_source = None

    history = []
    for g in range(generations):
        print(f"\n=== generation {g + 1}/{generations} ===")

        # Build the prompt
        prior_section = ""
        if prior_critique:
            prior_section = PRIOR_SECTION_TEMPLATE.format(
                stats=history[-1]["stats"],
                critique=prior_critique,
            )
        prompt = REWARD_GEN_TEMPLATE.format(
            task=task,
            state_spec=_format_state_spec(state_spec),
            prior_section=prior_section,
        )

        # Generate N candidates
        candidates = await generate_candidates(prompt, candidates_per_gen)

        # Train + eval each
        scored = []
        for i, src in enumerate(candidates):
            try:
                success, stats = train_and_eval(src)
                scored.append({"src": src, "success": success, "stats": stats})
                print(f"  candidate {i}: success={success:.3f}")
            except Exception as e:
                print(f"  candidate {i}: FAIL ({e})")

        if not scored:
            print("no candidates trained successfully; aborting")
            break

        # Pick best
        best = max(scored, key=lambda c: c["success"])
        print(f"  best success: {best['success']:.3f}")

        # Reflect
        critique = await reflect(best["src"], best["stats"], best["success"])
        print(f"  critique: {critique[:200]}...")

        # Persist
        gen_record = {
            "generation": g,
            "best_success": best["success"],
            "best_source": best["src"],
            "stats": best["stats"],
            "critique": critique,
        }
        (out_dir / f"gen-{g}.json").write_text(json.dumps(gen_record, indent=2))
        history.append(gen_record)

        prior_best_source = best["src"]
        prior_critique = critique

    return history


def _format_state_spec(spec: dict) -> str:
    return "\n".join(f"- {k}: {v}" for k, v in spec.items())