Train a tool-using agent with RL

This example trains Qwen3.5-4B with reinforcement learning (RL) on Loops to complete file tasks in a Linux sandbox. Each attempt gets its own Blaxel sandbox and two tools, shell and write_file. The model works for as many turns as it needs. A shell check the agent never sees then grades the files it left: reward 1 when the check passes, otherwise 0. The pattern fits any agent that acts on an environment through tool calls and is judged on the result: your tasks, your tools, and your grader, with the Loops cookbook running the sandboxes, sampling, and training.

Prerequisites

A Baseten workspace with Loops and its API key. Qwen3.5-4B trains only in the dedicated availability mode, which bills for GPU time.
A Blaxel workspace and API key.
uv and Python 3.12 or later.

Install the cookbook

Install the Loops cookbook from PyPI into a uv project. The cookbook needs Python 3.12 or later, installs the baseten-loops SDK, and provides the loops command. loops runs both recipes; each training config names its recipe with recipe = "rl" or recipe = "sft". Create a project first if you don’t have one:

uv init --python 3.12 loops-project
cd loops-project
uv add loops-cookbook

The package includes the cookbook’s examples. Run one using its bundled config path, such as examples/counting/train.toml, from your project’s directory. To change an example, copy its directory into your project at the same path. A config in your project wins over the installed one, and loops check prints the directory of the config it used. Create a .env file in your project’s directory and set BASETEN_API_KEY to a workspace API key. The commands load .env from the training config’s directory or the directory you run them from. Never commit .env. A run uses the serverless availability mode when your workspace has serverless access to the model, and the dedicated mode otherwise. To choose a mode, add BASETEN_LOOPS_SCOPE=serverless or BASETEN_LOOPS_SCOPE=deployment to .env. Install the Blaxel provider:

uv add "loops-cookbook[blaxel]"

Add your Blaxel credentials to .env as BL_WORKSPACE and BL_API_KEY. Without them, check stops with:

Blaxel needs BL_WORKSPACE and BL_API_KEY in the environment or in a .env beside the training config

Confirm your setup

Check the cookbook’s smallest pattern first. It needs no sandbox, so it confirms your API key and the model before you add Blaxel. From your project’s directory, run the following command:

uv run loops check patterns/custom_runtime/in_process_direct/train.toml

After the files the check used, the report confirms Loops serves the model at the configured sequence length, such as Loops model: Qwen/Qwen3.5-4B; max_seq_len 1024 (enabled up to <limit>). The report ends with:

Configuration and runtime checked; no sandbox or Loops trainer launched.

Read the tasks

examples/tool_agent/tasks.jsonl holds eight tasks. Each line has the files to create, an instruction for the agent, and a shell check that exits 0 when the task is done:

examples/tool_agent/tasks.jsonl

{"id": "port", "instruction": "In app.conf, change the port to 9090. Leave the other settings exactly as they are.", "files": {"app.conf": "name=shop\nport=8080\ndebug=true\n"}, "check": "[ \"$(cat app.conf)\" = \"$(printf 'name=shop\\nport=9090\\ndebug=true\\n')\" ]"}
{"id": "paid_orders", "instruction": "orders.csv lists orders, one per line after its header. Write to paid.txt how many orders have the status paid.", "files": {"orders.csv": "id,customer,status,amount\n1,ana,paid,30\n2,bo,unpaid,15\n3,cy,paid,12\n4,dee,refunded,40\n5,ed,unpaid,8\n6,flo,paid,25\n"}, "check": "[ \"$(tr -d ' \\n' < paid.txt)\" = 3 ]"}

The tasks form a ladder sized for the starting model. Easy ones, such as changing one setting, give early groups some successes. Harder ones each hold a trap a small model often falls into: grep paid also counts unpaid orders, so the obvious command gives 5 instead of 3. A group learns only when some attempts succeed and others fail, so these traps are where the learning signal comes from. See size tasks so attempts disagree.

Read the episode

examples/tool_agent/tool_episode.py defines ToolEpisode, one object per attempt. Its three methods are the whole environment. setup() writes the task’s files into /work in the attempt’s sandbox, before the agent starts:

examples/tool_agent/tool_episode.py

async def setup(self, setup: Setup) -> None:
    sandbox = setup.require_sandbox(Blaxel)
    files = self.task.value.files.items()
    box = await sandbox.client()
    await box.fs.write_tree(
        [{"path": path, "content": text} for path, text in files], WORKDIR
    )

run() is the agent. It sends the conversation and tools to the policy, runs each tool call in the sandbox, and appends the result, until the model replies without a tool call or reaches max_turns. It returns why it stopped, which the run reports as rollout/termination_reason/<reason>:

examples/tool_agent/tool_episode.py

for turn in range(1, self.options.max_turns + 1):
    try:
        reply = await conversation.generate(messages=messages, tools=TOOLS)
    except ContextWindowExceeded:  # the conversation outgrew the trajectory budget
        return ended("context_overflow", turn - 1, calls)
    # A cut reply trains with the episode's reward. To drop one the trajectory
    # budget cut (`reply.sample.clamped_from is not None`), call
    # `conversation.discard(reply)` before ending.
    if reply.truncated or reply.native_message is None:  # cut, or won't parse
        return ended("truncated" if reply.truncated else "unusable_reply", turn, calls)
    messages.append(reply.native_message)
    if not reply.tool_calls:
        return ended("completed", turn, calls)
    for (
        call
    ) in reply.tool_calls:  # in order: a later call may read an earlier one's file
        result = json.dumps(await self.call(sandbox, call))
        messages.append(
            {
                "role": "tool",
                "tool_call_id": call.id or "",
                "name": call.name or "",
                "content": result,
            }
        )
        calls += 1
return ended("max_turns", self.options.max_turns, calls)

The cookbook’s model client records the exact tokens the sampler returned for every generate() call. Those are the tokens that train. grade() stops the commands still running in the sandbox, then runs the task’s check. A process the agent detached isn’t stopped, so the check still trusts what the sandbox holds. See close the gaps a policy can exploit:

examples/tool_agent/tool_episode.py

async def grade(self, grading: Grading) -> float:
    sandbox = grading.require_execution_sandbox(Blaxel)
    box = await sandbox.client()
    # A command the agent deadline cut may still be running: stop what the sandbox
    # lists as running before the check. A job the agent detached is not stopped.
    for process in await box.process.list():
        if process.status == "running":
            await box.process.kill(process.pid)
    # Only a check that ran to the end passes: Blaxel's wait hands back a record
    # still ``running`` (exit code 0) when the sandbox stops answering mid-check.
    checked = await self.sh(sandbox, self.task.value.check)
    if checked["status"] == "completed" and checked["exit_code"] == 0:
        return 1.0
    await sandbox.ensure_alive()  # a failed check measures only a live sandbox
    return 0.0

A failed tool call, such as a command that exits non-zero or times out, goes back to the model as the tool’s result, and the episode continues. A failed call to Blaxel itself raises, and the cookbook retries the attempt in a fresh sandbox.

Read the training config

examples/tool_agent/train_blaxel.toml connects the episode, the tasks, and the model. These sections of it shape the run:

examples/tool_agent/train_blaxel.toml

[runtime.config]
episode = "tool_episode:ToolEpisode"
# After a delivered reply: graded from the sandbox state. Before one: unversioned.
agent_timeout_seconds = 180

# The agent's own limits, validated at check as its ToolOptions.

[runtime.config.options]
max_turns = 12
command_timeout_seconds = 20
max_output_chars = 2000

# One Blaxel sandbox per episode, provisioned before setup() and released by the run.

[runtime.config.sandbox]
provider = "blaxel"
memory_mb = 2048   # on Blaxel's base image, the default image
ttl = "1h"         # required: it must outlast one episode

[model]
base_model = "Qwen/Qwen3.5-4B"
lora_rank = 8
max_seq_len = 8192   # the whole conversation, tool results included
renderer_chat_template_kwargs = { enable_thinking = false }

[training]
mode = "sync"
group_size = 16
max_group_losses = 1   # a cloud sandbox can fail; its group still trains
groups_per_step = 1
max_steps = 5
max_consecutive_unusable_groups = 20   # stop if every group scores alike
learning_rate = 0.0001
max_steps_off_policy = 0
max_policy_delta_in_trajectory = 0

group_size = 16 runs each task 16 times per group. The advantage of each attempt is its reward minus the group’s mean.
max_group_losses = 1 lets a group still train when one sandbox fails.
max_steps = 5 keeps the first run short. Raise it once the per-group rewards show the tasks are sized well.
max_seq_len = 8192 bounds the whole conversation, tool results included.
ttl = "1h" is a backstop that deletes a sandbox a crashed run abandoned. The run releases each sandbox when its attempt ends.

Run the example

To check the configuration, run the following command. It validates the tasks and options and launches nothing:

uv run loops check examples/tool_agent/train_blaxel.toml

To train, run the following command. It starts Loops trainer and sampler deployments and one Blaxel sandbox per attempt, all of which bill while they run:

uv run loops run examples/tool_agent/train_blaxel.toml

The run writes its record under runs/tool-agent-blaxel/ in your project directory and logs the Loops run ID when it exits. The run ID is also in the run directory’s run.json.

Read the results

Read each update’s metrics in the run directory’s steps.jsonl:

Per-group rewards and discard/groups/no_variance. A task whose 16 attempts always score alike teaches nothing. Make it harder or easier.
rollout/termination_reason/*. A rising max_turns or truncated rate means the agent is running out of room rather than solving the task.
reward/mean and reward/mean_with_unusable. The first trains; the second includes groups the run rejected.

Check whether an RL run is learning explains each metric and what to do about it. Before you deploy a checkpoint, score it on tasks it never trained on with your evaluation.

Shut down the run

The dedicated trainer and sampler keep running after loops run exits. Deactivate the run when you finish to stop GPU billing without waiting for automatic shutdown. Run the following command with your run directory:

uv run loops deactivate "runs/tool-agent-blaxel/<run directory>"

The command prints the deactivated run ID. Saved checkpoints remain available.

Release sandboxes after a crash

Sandboxes bill until they’re deleted. If the run crashed, release any sandboxes it left with the following command:

uv run loops cleanup examples/tool_agent/train_blaxel.toml \
  --state-dir "runs/tool-agent-blaxel/<run directory>/state"

Adapt it to your tasks

Copy the example from the installed package into your project under a new name, such as my_agent:

COOKBOOK_DIR=$(uv run python -c 'import loops_cookbook; print(loops_cookbook.__path__[0])')
cp -R "$COOKBOOK_DIR/bundled/examples/tool_agent" my_agent

Then change three things in the copy:

Tasks: replace tasks.jsonl with your own files, instructions, and checks. Keep checks out of the instruction text.
Tools: change TOOLS and call() in tool_episode.py to the actions your agent takes in production.
Grader: change grade() to score what the task produces. Keep it independent of the agent’s own claims.

Then run uv run loops check my_agent/train_blaxel.toml, a short run, and read the per-group rewards before you scale up.

Next steps

Train a coding agent on Terminal-Bench: A Harbor agent on a real multi-turn benchmark.
Design RL environments and rewards: Task sizing, graders, and exploits.
Deploy a checkpoint: Serve the trained adapter.

Examples

Post-training

Models

Engines

Custom Docker servers

Custom Python models

Chains

Baseten for Model Labs

Train a tool-using agent with RL

Prerequisites

Install the cookbook

Confirm your setup

Read the tasks

Read the episode

Read the training config

Run the example

Read the results

Shut down the run

Release sandboxes after a crash

Adapt it to your tasks

Next steps

​Prerequisites

​Install the cookbook

​Confirm your setup

​Read the tasks

​Read the episode

​Read the training config

​Run the example

​Read the results

​Shut down the run

​Release sandboxes after a crash

​Adapt it to your tasks

​Next steps

Prerequisites

Install the cookbook

Confirm your setup

Read the tasks

Read the episode

Read the training config

Run the example

Read the results

Shut down the run

Release sandboxes after a crash

Adapt it to your tasks

Next steps