Train a vision model to count shapes

This example trains Qwen3.8-27B with reinforcement learning (RL) to count circles in a picture. Each picture has 9 to 20 circle outlines, and many of them overlap or sit inside another circle. Vision models count separate shapes well and overlapping shapes badly, so the base model starts with room to improve. Each episode is one model call: the policy gets one picture and one question, and replies with one number. The Loops cookbook draws the picture and grades the reply on your machine, so the run needs no sandboxes. Loops runs only the trainer and the sampler. The run takes about 25 minutes on serverless Loops. For an agent that takes many turns in a sandbox, see the coding agent example.

Prerequisites

A Baseten workspace with Loops access to Qwen/Qwen3.8-27B, and its API key. Qwen3.8-27B trains in either availability mode.
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 vision extra:

uv add "loops-cookbook[vision]"

The extra installs PyTorch and torchvision for the model’s image processor, which turns each picture into image tokens. They run on the CPU.

Run the example

To check the configuration, run the following command. It loads the episode and the task file, asks Loops about the model, and starts nothing:

uv run loops check examples/counting/train.toml

To train, run the following command. The run uses the serverless mode, billed for tokens, when your workspace has serverless access to the model. Otherwise, it uses the dedicated mode, billed for GPU time:

uv run loops run examples/counting/train.toml

The run writes to a new directory under runs/counting/. run.json has the Loops run ID, and steps.jsonl has the metrics of each update.

Read the episode

examples/counting/count_episode.py defines CountEpisode, the whole runtime of this example. The cookbook creates one episode for each attempt at a task, calls its run() method to get the reply, and calls its grade() method to get the reward. A task is a count and a seed. examples/counting/tasks.jsonl holds 512 tasks, and make_tasks.py writes them:

examples/counting/tasks.jsonl

{"name": "circles-000", "count": 15, "seed": 1806341205}

The seed sets where the circles go, so a task always draws the same picture. The attempts in a group see the same picture, and their rewards are comparable. run() draws the picture and sends it to the policy as an image part, beside the text of the question:

examples/counting/count_episode.py

    async def run(self, sample: Sampling) -> dict[str, JsonValue]:
        task = self.task.value
        picture = draw(task.count, task.seed)
        message = {
            "role": "user",
            "content": [
                {"type": "image", "image": picture},
                {"type": "text", "text": QUESTION},
            ],
        }
        reply = await sample.model.conversation().generate(messages=[message])
        self.guess = parse_count(reply.content)
        return {"count": task.count, "guess": self.guess}

draw() returns a 448 by 448 PNG. On Qwen3.8, that picture is 196 image tokens. The question asks for the number only. parse_count() keeps the first whole number in the reply. A reply with no whole number has no guess. The reward gives partial credit for a near miss:

examples/counting/count_episode.py

def score(guess: int | None, count: int) -> float:
    """1 for the right count, a quarter point less for each circle off, never below 0."""
    if guess is None:
        return 0.0
    return max(0.0, 1.0 - abs(guess - count) / 4)

A reply that is one circle off scores 0.75, and a reply that is four or more circles off scores 0. With partial credit, the attempts in a group seldom all score the same, so most groups have an advantage to train on.

Read the training config

examples/counting/train.toml names the episode and the task file:

examples/counting/train.toml

[runtime]
factory = "loops_cookbook.episodes:create_runtime"

[runtime.config]
episode = "count_episode:CountEpisode"
# Before the reply: unversioned. After it: graded, 0 if no answer was kept.
agent_timeout_seconds = 90

[runtime.files]
tasks = "tasks.jsonl"   # 512 pictures; written by make_tasks.py

The model and training settings:

examples/counting/train.toml

[model]
base_model = "Qwen/Qwen3.8-27B"
lora_rank = 8
max_seq_len = 1024
renderer_chat_template_kwargs = { enable_thinking = false }

[sampling]
max_generation_tokens = 16

[training]
mode = "sync"
group_size = 16
max_group_losses = 0
groups_per_step = 8   # eight pictures a step: each step trains on new ones
max_steps = 50
max_consecutive_unusable_groups = 20   # stop if every group scores alike
learning_rate = 0.00005
max_steps_off_policy = 0
max_policy_delta_in_trajectory = 0

Thinking is off, and a reply can have at most 16 tokens. A number needs only a few tokens.
Each update trains 8 pictures with 16 attempts each, for 50 updates. 50 updates use 400 of the 512 tasks, so every update trains on pictures the policy hasn’t seen.
mode = "sync" and the two 0 settings make every trained attempt come from the current policy. The trainer waits while the sampler counts.
max_consecutive_unusable_groups = 20 stops the run when 20 groups in a row have nothing to train, for example when every attempt scores the same.

Read the results

Monitor reward/mean in steps.jsonl. Each update trains on new pictures, so a reward that rises from update to update is a reward on pictures the policy hadn’t trained on. See Check whether an RL run is learning for the other metrics to watch. The cookbook’s recorded run of train.toml on serverless Loops took 23 minutes. The training reward rose from 0.76 in the first 5 updates to about 0.86 by update 10, and stayed near that value. The base model and the last checkpoint then each counted 64 held-out pictures, 8 times each:

	Base model	After 50 updates
Exact count	33%	60%
Off by one or less	69%	98%
Mean reward	0.72	0.90

The base model mostly counted one or two circles too many. After training, almost every reply was within one circle. A second run with another task order reached 48% exact counts after 10 updates and 52% after 20, so most of the gain comes early. To score your own checkpoint on pictures it didn’t train on, see Evaluate a post-trained model.

Shut down the run

loops run leaves the Loops run active. This serverless run incurs no token charges while idle, so deactivation is optional. To end it and its paired sampler, run the following command with your run directory:

uv run loops deactivate "runs/counting/<run directory>"

The command prints the deactivated run ID. Saved checkpoints remain available. See Deactivate a run for status checks and the Baseten CLI and REST API options.

Next steps

Vision: How image inputs become tokens, and which models accept them.
Fine-tune a vision model with the Loops SDK: SFT on images, written against the SDK directly.
Design RL environments and rewards: Graders, task sizing, and exploits.
Train a coding agent on Terminal-Bench: The same recipe with a multi-turn agent in Blaxel sandboxes.
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 vision model to count shapes

Prerequisites

Install the cookbook

Run the example

Read the episode

Read the training config

Read the results

Shut down the run

Next steps

​Prerequisites

​Install the cookbook

​Run the example

​Read the episode

​Read the training config

​Read the results

​Shut down the run

​Next steps

Prerequisites

Install the cookbook

Run the example

Read the episode

Read the training config

Read the results

Shut down the run

Next steps