Train a vision model to count shapes

Train Qwen3.6-35B-A3B with reinforcement learning to count overlapping circles in images on serverless Loops, with no sandboxes.

This example trains Qwen3.6-35B-A3B 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. The run takes about 24 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.6-35B-A3B, and its API key. Qwen3.6-35B-A3B 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 baseten-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 serverless supports the model and the run’s LoRA rank and sequence length, 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 "baseten-loops-cookbook[vision]"

The extra installs PyTorch and torchvision (CPU only) for the model’s image processor, which turns each picture into image tokens.

​

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

The check confirms the configuration and starts no trainer.

To train, run the following command. With serverless Loops, Baseten runs the trainer and sampler and bills for the tokens the run trains and samples, as long as serverless supports the model, the run’s low-rank adaptation (LoRA) rank, and its sequence length. Otherwise, the run uses dedicated capacity, which bills 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. 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"
agent_timeout_seconds = 90

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

agent_timeout_seconds = 90 caps the wait for a reply at 90 seconds. An attempt that returns no number scores 0.

The model and training settings:

examples/counting/train.toml

[model]
base_model = "Qwen/Qwen3.6-35B-A3B"
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
max_steps = 50
max_consecutive_unusable_groups = 20
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", max_steps_off_policy = 0, and max_policy_delta_in_trajectory = 0 make every trained attempt come from the current policy. The trainer waits while the sampler counts.
• max_group_losses = 0 trains a group only when every attempt in it finished.
• 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. For the other metrics to watch, see Check whether an RL run is learning.

The cookbook’s recorded run of train.toml on serverless Loops took 24 minutes. The base model and the last checkpoint then each counted 64 held-out pictures, 8 times each:

	Base model	After 50 updates
Exact count	26%	66%
Off by one or less	62%	93%
Mean reward	0.64	0.90

The base model mostly counted one or two circles too many. After training, almost every reply was within two circles, and most misses were one too few. Every picture size gained: exact counts rose from 43% to 83% at 9 to 12 circles, and from 12% to 53% at 17 to 20.

To score your own checkpoint on pictures it didn’t train on, see Evaluate a post-trained model.

​

Resume a stopped run

If the run stops before its last update, continue it with --resume. The command finds the run of this config that trained most recently, and continues from the last checkpoint it saved:

uv run loops run examples/counting/train.toml --resume

max_steps counts the updates of the whole run, so the resumed run makes only the updates that are left. To check the resumed run first, run uv run loops check examples/counting/train.toml --resume. If no run has trained, or the run saved no checkpoint, the command exits with code 2 and starts nothing.

​

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. For status checks and the Baseten CLI and REST API options, see Deactivate a run.

​

Deploy the checkpoint

The run’s sampler checkpoint is ready to serve. The deployment loads the adapter on the base weights, downloads the base weights with the hf_access_token in workspace secrets, and bills for its GPU while it’s live. With the Baseten CLI installed, list the run’s checkpoints and deploy the one you want:

baseten loops checkpoint list --run-id "<run_id>"
baseten loops checkpoint deploy --checkpoint-id "<checkpoint_id>"

The deploy command prompts for a model name, GPU type, GPU count, and the Hugging Face secret name. Deploy a checkpoint covers the REST API and dashboard paths.

​

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.