Training environment

The environment is what the policy acts in and what produces the observations a reward scores. Its shape depends on the task:

Environment	Example	Each episode is
A dataset of verifiable problems	Math word problems with known answers	One prompt, and the completions the policy samples for it
A simulator	A MuJoCo or Brax physics model	A trajectory of observations, actions, and rewards
A sandboxed world	A filesystem and tools an agent explores	A multi-turn trajectory of tool calls and their results

Whatever the shape, the environment has to be cheap to reset, isolated between episodes, and reproducible, so that a retried episode produces the same result.

Use a dataset as the environment

For RL on reasoning tasks, the environment is a set of prompts with checkable answers. The driver samples a batch each iteration, and the reward compares each completion with the answer. Load the dataset in a task with cache="auto", or bundle a small one with the code. The loop on these pages rotates through a fixed list of prompts:

        
    
def _sample_prompts(iteration: int) -> list[tuple[str, str]]:
    start = (iteration * PROMPTS_PER_ITER) % len(DATASET)
    return [DATASET[(start + i) % len(DATASET)] for i in range(PROMPTS_PER_ITER)]

Sampling deterministically from the iteration number means a resumed loop picks up the same prompts it would have used.

An agentic environment is a world the policy explores, such as a directory of documents, a repository, or a database. Building it can take longer than an episode, so build each world once in a cached task and pass it to every episode:

agentic_rl_durable.py
                
@sandbox_env.task(cache="auto")
async def build_world(world_id: str) -> ROVolume:
    """Populate the world's filesystem once; sealed as a read-only volume every trial forks."""
    secret_doc, secret = _world_secret(world_id)
    vol = Volume.new(name=f"world-{world_id}-{flyte.ctx().action.run_name}")
    await vol.mount()
    root = Path(vol.mount_path) / "docs"
    root.mkdir(parents=True)
    for i in range(N_DOCS):
        lines = [f"title: report {i} for {world_id}", f"owner: team-{i % 3}", f"pages: {10 + i}"]
        if i == secret_doc:
            lines.append(f"secret_sha256: {hashlib.sha256(secret.encode()).hexdigest()}")
        (root / f"doc_{i}.txt").write_text("\n".join(lines) + "\n")
    return await vol.finalize(message=f"world {world_id}")

The world is a read-only volume. Each episode forks it into an instant private copy and mounts the fork:

agentic_rl_durable.py
                
                    
                
            
                
            

                
                    
                
            
@sandbox_env.task(retries=1, timeout=flyte.Timeout(max_runtime=300))
async def generate(spec: TrialSpec, weights: Weights, world: ROVolume, flaky_rate: float = 0.0) -> Trajectory:
    """One trial: fork the world, run the agent loop against it, return what happened.

    Idempotent by construction — re-running produces the same trajectory (seeded by trial_id),
    so a retry is a clean regeneration, not a replay of stale sandbox state.
    """
    t0 = time.monotonic()
    attempt = int(os.environ.get("FLYTE_ATTEMPT_NUMBER", "0"))
    rng = random.Random(_seed("trial", spec.trial_id))
    if attempt == 0 and rng.random() < flaky_rate:
        raise RuntimeError(f"simulated sandbox failure for {spec.trial_id} (attempt 0)")  # -> retried

    # [S2] instant private copy of the world for this trial
    fork_name = f"{world.name}-{spec.trial_id}"
    ws = await world.fork(name=fork_name)
    await ws.mount()
    docs = Path(ws.mount_path) / "docs"

    # the "agent": pick documents to read in an order driven by the policy's prior for this world
    prior = list(weights.doc_prior.get(spec.world_id) or [1.0] * N_DOCS)
    order: list[int] = []
    remaining = list(range(N_DOCS))
    while remaining:
        w = [prior[i] for i in remaining]
        pick = rng.choices(remaining, weights=w, k=1)[0]
        order.append(pick)
        remaining.remove(pick)

    turns, answer = [], None
    for doc in order[: spec.max_turns]:
        text = (docs / f"doc_{doc}.txt").read_text()
        found = "secret_sha256:" in text
        turns.append({"doc": doc, "found": found})
        await asyncio.sleep(0.2)  # a tool call
        if found:
            answer = text.split("secret_sha256:")[1].strip()
            break

    written = []
    (docs.parent / "answer.txt").write_text(answer or "")
    written.append("answer.txt")
    if rng.random() < 0.3:  # untidy agents leave scratch files behind (the judge cares)
        (docs.parent / "scratch.tmp").write_text("notes")
        written.append("scratch.tmp")

    return Trajectory(
        trial_id=spec.trial_id,
        world_id=spec.world_id,
        weights_version=weights.version,
        turns=turns,
        answer=answer,
        files_written=written,
        elapsed_s=round(time.monotonic() - t0, 2),
        fork_name=fork_name,
    )

cache="auto" on the build task means every run, including one that re-scores old episodes with a new reward, gets the same world. That keeps the cache keys of the downstream tasks stable.

Run each episode in its own task

Running an episode as its own task gives it the properties an environment needs:

Isolation. Each episode gets a fresh container, so a policy that deletes files or fills the disk only affects its own episode.
Limits. timeout=flyte.Timeout(max_runtime=300) stops an episode that loops forever, and retries=1 reruns one that hits a transient failure.
Reproducibility. Seed the episode’s randomness from its ID, so a retry produces the same trajectory instead of a different one.
Separate sizing. The environment tier gets CPU and memory for the simulator or tools, with no GPU.

Episodes are independent, so the driver launches all of them at once with asyncio.create_task and collects them with asyncio.as_completed. It awaits each episode in a try block, so an episode that still fails after its retries is skipped instead of failing the iteration. See Training loop.

Run model-written code safely

When the policy writes code that runs during an episode or during scoring, such as a generated function and the test suite that checks it, run that code in a sandbox rather than in the task’s own process. See Code sandboxing, and the MLE bot tutorial for an agent whose generated code runs in a sandbox.

Serve the environment as an app

Some environments are long-running servers that clients step through over HTTP, such as an environment that keeps game state between calls. Deploy the server as an app and connect to it from the rollout tasks by its URL. The app scales and restarts independently of the training loop, and any number of rollout tasks can share it.

Reward setup: score what happened in each episode.
Inference rollouts: generate the policy’s actions.
Durable agentic RL with SkyRL and Harbor: the world-building and episode code on this page, in a complete loop.

Training environment

Use a dataset as the environment

Build a world once and share it

Run each episode in its own task

Run model-written code safely

Serve the environment as an app

Related