Guides / Python / RL rollouts and evals

RL rollouts and evals

An eval or an RL run needs many copies of one environment, each sealed from the others and from your machine, reset cheaply between episodes. A kosa8 sandbox is that copy; a snapshot is the reset.

The shape of it

1. Prepare once. Make a sandbox, install the task's dependencies, copy in the repo. Save it as a template or keep it as a snapshot.
2. Fan out. fork(n) makes n running copies of that state.
3. Roll out. Each policy or model works in its own copy through the toolbox.
4. Score. Run the hidden tests in the copy and read the exit code.
5. Throw away. Delete the copies. The prepared state is untouched.

1. Prepare the environment

from kosa8 import Kosa8Client

client = Kosa8Client()
base = client.create_sandbox(name="task-base", egress_allow=["pypi.org", "files.pythonhosted.org"])
sb = client.sandbox(base["id"])
sb.git("clone", url="https://github.com/pallets/click", path="repo")
sb.shell("cd /workspace/repo && pip install -q -e . pytest")

For an environment you will reuse across runs, put the setup in a Dockerfile and build a template once: client.build_template("click-env", "./env"). Then every sandbox starts from it with create_sandbox(template="click-env"), with no install step.

2. Fork one rollout per sample

f = client.fork(base["id"], 8, prefix="ep-")
episodes = [x["id"] for x in f["forks"]]      # ep-1 … ep-8, all running, all identical

Forking snapshots the base once and restores it eight times. The base keeps running. Each fork has its own disk, memory and network rules, so one episode cannot see or spoil another.

3. Run the policy in each copy

Give each episode its own tools. With the OpenAI Agents SDK (see agent frameworks):

from concurrent.futures import ThreadPoolExecutor
from agents import Agent, Runner
from kosa8.integrations.openai_agents import sandbox_tools

TASK = "Fix the failing test in /workspace/repo. Do not edit the tests."

def rollout(sid):
    agent = Agent(name="solver", instructions=TASK, tools=sandbox_tools(client, sid))
    return Runner.run_sync(agent, TASK, max_turns=30).final_output

with ThreadPoolExecutor(8) as pool:
    transcripts = list(pool.map(rollout, episodes))

The client is safe to share between threads; each call is its own request.

4. Score with a reward the agent cannot see

Copy the grader in only after the rollout, so the policy cannot read or edit it, and score on the exit code:

GRADER = open("hidden_tests.py").read()

def reward(sid):
    sb = client.sandbox(sid)
    sb.write_file("/workspace/.grader/test_hidden.py", GRADER)
    r = sb.shell("cd /workspace/repo && python -m pytest -q /workspace/.grader", timeout_s=300)
    return 1.0 if r.ok else 0.0

rewards = [reward(sid) for sid in episodes]

For a partial score, have the grader print JSON and parse r.stdout. With run_code you can score in Python directly and get back charts as PNG bytes in r.artifacts.

5. Clean up

for sid in episodes:
    client.delete_sandbox(sid)

Or let kosa8 do it: give the base ttl_min= or auto_stop_min=, and make throwaway sandboxes ephemeral=True, so they are deleted the moment they stop, even if your script crashes.

Branching inside an episode

For search (tree search, best-of-N per step, undoing a bad action), take a checkpoint before a risky step and rewind if it fails. The rest of the episode is unaffected:

client.checkpoint("ep-1", label="before migration")
r = client.shell("ep-1", "cd /workspace/repo && ./migrate.sh")
if not r.ok:
    client.rewind("ep-1")          # back to the checkpoint: disk and memory

To try several next actions at once and keep the copy where one worked, fork the episode once per action, run each in its own copy, and keep the first that passes:

patches = ["patch-a.diff", "patch-b.diff", "patch-c.diff"]
tries = [x["id"] for x in client.fork("ep-1", len(patches), prefix="ep-1-try")["forks"]]
with ThreadPoolExecutor(len(tries)) as pool:
    oks = list(pool.map(lambda t: client.shell(
        t[0], f"cd /workspace/repo && git apply /workspace/{t[1]} && python -m pytest -q").ok,
        zip(tries, patches)))
winner = next((sid for sid, ok in zip(tries, oks) if ok), None)
for sid in tries:
    if sid != winner:
        client.delete_sandbox(sid)

explore() does this in one call for commands that need only an image, not the workspace: each variant runs in a fresh container of image, and it keeps the fastest copy that exits 0. See explore in depth.

Evals over a dataset

The same pattern runs a benchmark: one forked sandbox per task, the model's answer applied, the task's tests run, the pass rate counted.

results = []
for task in dataset:                               # each: {"id", "setup", "check"}
    sid = client.fork(base["id"], 1, prefix=f"eval-{task['id']}-")["forks"][0]["id"]
    try:
        client.shell(sid, task["setup"])
        answer = model_solve(task, sandbox_tools(client, sid))
        results.append(client.shell(sid, task["check"]).ok)
    finally:
        client.delete_sandbox(sid)
print(f"pass rate {sum(results)}/{len(results)}")

Things to watch

• Network: sandboxes start with the network closed. Open only what setup needs, and consider a separate base for rollouts with egress denied, so a policy cannot fetch answers.
• Resources: every fork is a microVM with its own memory. Size mem_mb on the base, and keep the number of live forks to what the machine holds.
• Time: pass timeout_s to shell and run_code, so a hung episode ends with timed_out=True rather than blocking the batch.
• Audit: every command an agent ran is in the kosa8 audit log, which helps when a reward looks too good.

See also: explore and fork in depth and explore.