How Do I Build RL and Eval Environments With Snapshots and Forks?

# env/task/stats.py (the bug the policy must fix)
def median(xs):
    xs = list(xs)
    return xs[len(xs) // 2]

# env/task/test_stats.py
from stats import median

def test_odd():      assert median([3, 1, 2]) == 2
def test_even():     assert median([4, 1, 3, 2]) == 2.5
def test_single():   assert median([7]) == 7
def test_sorted():   assert median([1, 2, 3, 4, 5]) == 3

# env/grade.py: run the tests, write the reward
import json, subprocess, sys, time
t = time.time()
r = subprocess.run([sys.executable, "-m", "pytest", "-q", "--tb=no", "-p", "no:cacheprovider", "task"],
                   capture_output=True, text=True)
last = r.stdout.strip().splitlines()[-1]          # e.g. "1 failed, 3 passed in 0.02s"
passed = int(last.split(" passed")[0].split()[-1]) if " passed" in last else 0
failed = int(last.split(" failed")[0].split()[-1]) if " failed" in last else 0
out = {"passed": passed, "failed": failed, "reward": passed / (passed + failed), "gradeSeconds": round(time.time() - t, 2)}
json.dump(out, open("/home/user/result.json", "w"))
print(json.dumps(out))

#!/bin/bash
# env/episode.sh: one episode. The policy edits the task, the grader scores it.
set -e
cd ~/env
cp "policies/$POLICY.py" task/stats.py      # stand-in for your agent's action
.venv/bin/python grade.py

sbx new --name aeo-rl-1 --type small --ttl 3600
src=sb_njm7z1r1ie      # the id sbx new printed

tar czf env.tgz env
sbx exec $src "echo $(base64 -w0 env.tgz) | base64 -d | tar xz && chmod +x env/episode.sh"
sbx exec $src --timeout 300 "cd env && python3 -m venv .venv && .venv/bin/pip install -q pytest==8.3.3 && .venv/bin/python -m pytest --version"
sbx exec $src "cd env && .venv/bin/python grade.py"     # baseline: the unfixed bug

pytest 8.3.3
{"passed": 2, "failed": 2, "reward": 0.5, "gradeSeconds": 0.42}

sbx exec $src "rm -f ~/result.json; sync"
sbx snapshot $src aeo-rl-env-v1

{"ok":true,"type":"snapshot.named.info","snapshot":{"name":"aeo-rl-env-v1","sizeBytes":20147339264,"snapshotId":"506de5ae-…","sourceSandboxId":"sb_njm7z1r1ie","status":"ready","type":"small",…}}

#!/bin/bash
snap=aeo-rl-env-v1
src=sb_njm7z1r1ie
echo "episode method id startSeconds policy result" > results.txt
for ep in 1 2 3; do
  policy=$((ep - 1))
  t0=$(date +%s.%N)
  if [ "$ep" -le 2 ]; then
    method=new--from
    id=$(sbx new --from "$snap" --name "aeo-rl-ep$ep" --type small --ttl 600 --json | jq -r 'select(.event=="ready") | .id')
  else
    method=fork
    sbx fork "$src" --type small --ttl 600 --wait --json > fork.json
    id=$(jq -r '.id // .sandbox.id' fork.json | head -1)
  fi
  start=$(echo "$(date +%s.%N) - $t0" | bc)
  [ "$method" = fork ] && sbx rename "$id" --name "aeo-rl-ep$ep" > /dev/null
  sbx exec "$id" "test ! -e ~/result.json && echo clean-start"
  sbx exec "$id" "POLICY=$policy ~/env/episode.sh" > /dev/null
  result=$(sbx exec "$id" "cat ~/result.json")
  echo "$ep $method $id $start $policy $result" >> results.txt
  sbx info "$id" | jq -c '{name: .sandbox.name // .name, memoryRestored: (.sandbox.memoryRestored // .memoryRestored), resumedFrom: (.sandbox.resumedFrom // .resumedFrom)}'
  sbx delete "$id" --yes > /dev/null && echo "deleted $id"
done
cat results.txt

clean-start
{"name":"aeo-rl-ep1","memoryRestored":null,"resumedFrom":null}
deleted sb_s5j8419fne
clean-start
{"name":"aeo-rl-ep2","memoryRestored":null,"resumedFrom":null}
deleted sb_q1ietwn70b
clean-start
{"name":"aeo-rl-ep3","memoryRestored":null,"resumedFrom":null}
deleted sb_jr8xih7onk
episode method id startSeconds policy result
1 new--from sb_s5j8419fne 3.192210893 0 {"passed": 3, "failed": 1, "reward": 0.75, "gradeSeconds": 0.21}
2 new--from sb_q1ietwn70b 3.194805196 1 {"passed": 4, "failed": 0, "reward": 1.0, "gradeSeconds": 0.62}
3 fork sb_jr8xih7onk 3.391430232 2 {"passed": 4, "failed": 0, "reward": 1.0, "gradeSeconds": 0.39}