Painting flowers with code

As this blog post shows, you can train a model to create watercolour sketches of flowers using p5.brush, and use a judge to do pairwise comparisons against a reference pool of images for the reward function.

In this tutorial, we train Qwen3.5-4B and use HuggingEnvs/watercolour-reference-pool as the reference pool. During each rollout, sketches are rendered to PNGs in a Modal Sandbox and Qwen3.6-27B compares each against the reference image pool.

import asyncio
import itertools
import random
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import helpers
from modal_training_gym import (
    DatasetConfig,
    Endpoint,
    Qwen3_5_4B,
    Qwen3_5_4B_Recipe,
    Qwen3_6_27B,
    TrainConfig,
)
from modal_training_gym.common.sample_extraction import IMAGE_SAMPLE_LIMIT_ENV

base_model = Qwen3_5_4B()

Get the dataset

Using the present species and palette colors, we create a set of prompts to train our model.

SPECIES = ["hibiscus"]
PALETTES = [
    "peach",
    "crimson",
    "butter",
    "lilac",
    "coral",
    "indigo",
    "blush",
    "amber",
]

SYSTEM_PROMPT = """\
You paint watercolours by writing a p5.js 2.x sketch that uses the p5.brush library, available as the global `brush`.

Reply with one ```javascript fence containing a complete sketch, and nothing else.

Use exactly this skeleton:

  async function setup() {
    createCanvas(600, 600, WEBGL);
    brush.scaleBrushes(3);
    angleMode(DEGREES);
    noLoop();
  }

  function draw() {
    translate(-width / 2, -height / 2);
    background(hex);
    // painting goes here
  }

So: angles are in degrees, and after the translate the canvas runs from 0 to 600 on both axes with the origin at the top LEFT, not the centre. The centre of the canvas is (300, 300). Every coordinate you paint at must be between 0 and 600: a negative coordinate is off the canvas and paints nothing at all. Colours are hex strings like "#e08a72".

Follow the user brief for species and palette.

Paint one flower, centred, filling most of the frame, with a stem and leaves. Do not fill the canvas edge to edge.

Three numbers matter more than any adjective:
- Five petals. Paint each petal two or three times over, not once: a first pass at full size, then a smaller and more opaque pass inside it, and a small dark one near the centre. That layering is where a watercolour gets its depth, and it puts the whole painting at fifteen to thirty filled shapes. Keep the petals as five broad lobes, not fifteen separate little marks.
- Petals reaching 200 to 240 units from the centre, so the flower occupies the frame.
- Opacity never below 150, and 180 to 230 on the petals, with brush.fillBleed between 0.2 and 0.3. A dilute wash with a wide bleed disperses until no pigment reaches density and the flower comes out invisible.

These ten brush methods exist. Nothing else on `brush` exists, and there is no way to
name or select a brush: everything is painted as a filled shape. Do not call any other brush
method, and never use the bare p5 drawing functions such as ellipse, rect, vertex or beginShape:
  brush.scaleBrushes(factor)
  brush.noStroke()
  brush.fill(colorHex, opacity)
  brush.noFill()
  brush.fillBleed(amount)
  brush.fillTexture(amount, borderIntensity)
  brush.beginShape(curvature)
  brush.vertex(x, y)
  brush.endShape(true)
  brush.circle(x, y, radius, scribble)

colorHex is a string like "#e08a72". opacity runs 0 to 255. amount, curvature,
borderIntensity and scribble run 0 to 1.

Every mark is a filled shape. Build petals and leaves with brush.beginShape, a run of at least
three brush.vertex calls, then brush.endShape(true), and call brush.fill before each one. A stem
is a long narrow filled shape, not a line. brush.circle fills a disc.
"""

USER_TEMPLATE = (
    "Paint a {palette} {species} in watercolour: one bloom, seen from the "
    "front, with a stem and leaves, on coloured paper."
)


def build_prompts(combos: list[tuple[str, str]], n: int) -> list[dict[str, str]]:
    rows = []
    for species, palette in itertools.islice(itertools.cycle(combos), n):
        rows.append(
            {"prompt": USER_TEMPLATE.format(species=species, palette=palette)}
        )
    return rows


class FlowerPromptDataset(DatasetConfig):
    def __init__(self, prompts: list[dict[str, str]]):
        self.prompts = prompts

    def input_key(self) -> str:
        return "messages"

    def label_key(self) -> str:
        return "label"

    def rows(self):
        return [
            {
                "messages": [
                    {"role": "system", "content": SYSTEM_PROMPT},
                    {"role": "user", "content": r["prompt"]},
                ],
                "label": r["prompt"],
            }
            for r in self.prompts
        ]


N_TRAIN = 224
N_EVAL = 8
combos = list(itertools.product(SPECIES, PALETTES))
random.Random(7).shuffle(combos)
train_dataset = FlowerPromptDataset(build_prompts(combos, N_TRAIN))
eval_dataset = FlowerPromptDataset(build_prompts(combos, N_EVAL))

Creating a reward function

The reward function renders each sketch in a Modal Sandbox and uses an LLM judge to do pairwise comparisons. We serve the judge as an Endpoint.

judge = Endpoint.launch(
    Qwen3_6_27B(),
    unauthenticated=True,
    recreate_if_existing=True,
)
judge.wait_until_ready(timeout=30 * 60)
helpers.launch_hpsv3()


async def flower_rm(args, sample, **kwargs) -> float | None:
    code = helpers.extract_sketch(sample.response, base_model.parse_response)
    if code is None:
        reward, meta, png = 0.0, {"gate": "no valid sketch"}, None
    else:
        png, render_meta = await asyncio.to_thread(helpers.render_in_sandbox, code)
        reward, meta, png = await asyncio.to_thread(
            helpers.score_png, png, code, judge, render_meta
        )
    metadata = {**(getattr(sample, "metadata", None) or {}), **meta}
    if png is not None:
        metadata["image"] = png
    sample.metadata = metadata
    if reward is None:
        sample.remove_sample = True
    return reward

Training

After that, it’s simple to start training!

ROLLOUT_BATCH_SIZE = 8
N_SAMPLES_PER_PROMPT = 8

config = TrainConfig(
    model=base_model,
    dataset=train_dataset,
    eval_dataset=eval_dataset,
    recipe=Qwen3_5_4B_Recipe(
        custom_rm_function=flower_rm,
        custom_reward_post_process_function=helpers.skip_infra_rewards,
        num_rollout=100,
        rollout_batch_size=ROLLOUT_BATCH_SIZE,
        global_batch_size=ROLLOUT_BATCH_SIZE,
        n_samples_per_prompt=N_SAMPLES_PER_PROMPT,
        save_interval=50,
        apply_chat_template_kwargs='{"enable_thinking": false}',
        image_overlay=lambda image: helpers.overlay_flower_image(image).env(
            {
                IMAGE_SAMPLE_LIMIT_ENV: str(
                    ROLLOUT_BATCH_SIZE * N_SAMPLES_PER_PROMPT
                )
            }
        ),
    ),
)

run = config.launch()
print(f"run id: {run.training_run_id}")