プロシージャル水彩画生成LLMの実装：TRLとOpenEnvによるコード実行環境とGRPO強化学習

TL;DR

• Hugging Face TRLのGRPO（Group Relative Policy Optimization）と外部実行環境OpenEnvを統合し、LLMにp5.brushを用いたJavaScript水彩画描画コードを生成させる強化学習パイプラインを構築。
• 構文成否、HPS（美的スコア）、VLMペアワイズ比較を組み合わせた3層ハイブリッド報酬関数により、Diffusionモデルに依存しない解像度非依存のプロシージャルアート生成を実現。
• 価値モデル（Critic）を廃したGRPOとLoRA（PEFT）の併用により、8BクラスLLMの強化学習を単一ノード（RTX 4090 / 24GB VRAM）で安定実行するためのハイパーパラメータとKLダイバージェンス制御手法を提示。

背景と技術的課題

Latent Diffusionモデル（SDXLやFLUX）は高品質な画像生成を実現する一方、出力はラスタライズされたピクセルグリッドに固定される。解像度非依存のベクターアセット取得、ストローク単位での描画ステップの編集・再構成、ペンプロッタやCADといった実機出力との連携が必要なユースケースにおいて、このピクセル固定性は根本的なアーキテクチャ上の制約となる。

この制約を打開する手法として、LLMにSVGやp5.js等の描画コードを直接生成させるプロシージャルアプローチがある。しかし、従来のSFT（教師あり微調整）モデルは「文法的に有効なJavaScript/Canvasコード」を出力できても、その描画結果が「視覚的に優れた水彩画」となる品質保証が困難であった。コード生成モデルの重みに美的基準を直接獲得させるには、生成コードのレンダリング結果（画像）に対する報酬を用いた強化学習（RLHF/RLAIF）の適用が不可欠となる。

従来、言語モデルの強化学習にはPPO（Proximal Policy Optimization）が主流であったが、方策モデル（Actor）、参照モデル（Reference）、報酬モデル（Reward）、価値モデル（Critic）を同時にVRAM上に展開する必要があり、8Bクラスのモデルでも最低32GB以上のメモリが要求されていた。本構成では、Criticモデルを廃しグループサンプリング内の相対報酬（Z-score正規化）でアドバンテージを算出するGRPOを採用。コード実行を外部の分離サンドボックス（OpenEnv）に委譲することで、コンシューマ向け24GB VRAM環境でのエンドツーエンド強化学習を可能にした。

コアアーキテクチャ・仕組みの解説

本パイプラインは、非同期に協調する生成・実行・評価・方策更新の4フェーズで構成される。

[LLM (Actor Policy: LLaMA-3-8B + LoRA)] --(1. Generate JS Code)--> [OpenEnv (Node.js/headless-gl)]
              ^                                                                  |
              |                                                           (2. Render Image)
      (4. GRPO Update)                                                           |
              |                                                                  v
      [KL Divergence] <-- [3-Tier Hybrid Reward] <----------------------- [Image Buffer]
                          |-- r_exec  : Syntax & Runtime Success Check
                          |-- r_aes   : HPSv2-based Aesthetic Score
                          `-- r_style : VLM Pairwise Tournament vs References

アーキテクチャの成否を分ける設計上の要所は、「LLMのアクション空間の適切な制約」と「多層化された報酬のゲーティング」にある。

1. アクション空間の制約（p5.brush DSLの定義）

汎用のJavaScriptコードを自由に生成させると、探索空間が膨大化し、DOMアクセスや無意味な計算ループによる学習停滞を招く。そのため、p5.brushライブラリが備えるAPI群から、水彩表現に特化した10種類のプリミティブ（brush.scaleBrushes, brush.fillBleed, brush.fillTexture, brush.bleed, brush.stroke等）のみを使用可能な構文制約としてプロンプトで定義する。さらに、推論時のLogit Biasを併用して不要な制御構文の出現確率を抑制し、LLMのアクション空間を実質的なドメイン固有言語（DSL）に限定する。

2. 3層ハイブリッド報酬関数の設計

生成された描画コードはOpenEnvサンドボックスへ送られ、Node.js環境上のheadless-glにより即座にオフスクリーンレンダリングされる。得られた実行結果および画像バッファに対し、以下の3段階で報酬を算出する。

1. 実行成否報酬（r_exec）: 構文エラー、タイムアウト（2.0秒超過）、またはレンダリング例外が発生した場合は -1.0 を即座に返却。以降の重い視覚評価をスキップ（ショートサーキット評価）する。正常描画時はベースラインとして +0.1 を付与。
2. 絶対的美的スコア（r_aes）: レンダリング画像に対し、HPSv2等の美的品質評価モデルを用いて構図・色彩調和スコアを算出（[-1.0, 1.0] にスケーリング）。
3. スタイル一致度ペアワイズ評価（r_style）: 高品質な手描き水彩画のリファレンスプールから画像を抽出し、VLMに対して「どちらが水彩特有の滲み・紙テクスチャ・エッジの揺らぎを精緻に表現しているか」を判定させる。ペアワイズ勝敗結果をBradley-Terryモデルにより相対効用値へ変換する。

総合報酬 $R$ は、以下の通り定義される。

R = r_exec + (0.5 * r_aes) + (0.4 * r_style)  [ただし r_exec < 0 の場合は R = r_exec]

実践ハンズオン：動く実装コードとパイプライン

以下は、Hugging Face trl（v0.12以上）およびpeftを用いたGRPOTrainerの実装スクリプトである。実行には Python 3.10+、torch >= 2.4.0、およびローカルで稼働するOpenEnvコンテナ（ポート8080）を前提とする。

import os
import torch
from datasets import load_dataset
from peft import LoraConfig
from transformers import AutoModelForCausalLM, AutoTokenizer
from trl import GRPOConfig, GRPOTrainer

# 1. サンドボックス実行環境のクライアント初期化
# OpenEnv: Node.js + p5.js + headless-gl を内包する分離コンテナ
from openenv import OpenEnvClient
env = OpenEnvClient(sandbox_url="http://localhost:8080")

model_id = "meta-llama/Meta-Llama-3-8B-Instruct"
tokenizer = AutoTokenizer.from_pretrained(model_id)
if tokenizer.pad_token is None:
    tokenizer.pad_token = tokenizer.eos_token

# 2. VLMペアワイズ評価のインターフェース
def evaluate_watercolor_style_via_vlm(image_bytes: bytes, prompt: str) -> float:
    """
    リファレンス画像とのペアワイズ比較をVLM API経由で実行。
    水彩画特有の滲みやテクスチャの忠実度をスコア化（-1.0 〜 1.0）。
    """
    # 外部VLM API（Claude / GPT-4o等）またはローカル軽量VLMを呼び出す
    # 実装例: 勝敗確率に基づくBradley-Terry効用ロジットの算出
    return 0.2  # プレースホルダ

# 3. TRL準拠の複合報酬関数
def procedural_watercolor_reward_func(completions, prompts=None, **kwargs):
    rewards = []
    for code in completions:
        # Step 1: OpenEnvサンドボックスによる安全なコード実行
        exec_result = env.run_code(code, timeout_sec=2.0)
        
        # 実行時エラーまたは画像生成失敗時は即座にペナルティ
        if not exec_result.success or exec_result.image_bytes is None:
            rewards.append(-1.0)
            continue
            
        # Step 2: HPS（美的スコア）による構図・配色の絶対評価
        hps_score = env.compute_aesthetic_score(
            exec_result.image_bytes,
            prompt="high quality watercolor painting, natural bleeds and texture"
        )
        
        # Step 3: VLMによるスタイル一致度ペアワイズ評価
        style_score = evaluate_watercolor_style_via_vlm(
            exec_result.image_bytes,
            prompt="watercolor style fidelity"
        )
        
        # 総合報酬の統合（正常終了ベースライン +0.1 を付与）
        total_reward = 0.1 + (0.5 * hps_score) + (0.4 * style_score)
        rewards.append(total_reward)
        
    return rewards

# 4. LoRA（PEFT）構成：24GB VRAMでの学習を成立させる必須設定
peft_config = LoraConfig(
    r=16,
    lora_alpha=32,
    target_modules=[
        "q_proj", "k_proj", "v_proj", "o_proj",
        "gate_proj", "up_proj", "down_proj"
    ],
    lora_dropout=0.05,
    bias="none",
    task_type="CAUSAL_LM",
)

# 5. GRPOハイパーパラメータ設定
training_args = GRPOConfig(
    output_dir="./grpo-watercolor-llama3-8b",
    learning_rate=1e-5,
    lr_scheduler_type="cosine",
    beta=0.04,                     # KLペナルティ係数（モード崩壊の抑制）
    num_generations=8,             # プロンプトあたりのグループ生成数 (G)
    per_device_train_batch_size=1,
    gradient_accumulation_steps=4,
    max_prompt_length=512,
    max_completion_length=1024,
    logging_steps=10,
    save_steps=100,
    bf16=True,
    report_to="none"
)

# 6. モデルのロードとTrainer構築
model = AutoModelForCausalLM.from_pretrained(
    model_id,
    torch_dtype=torch.bfloat16,
    device_map="auto"
)

trainer = GRPOTrainer(
    model=model,
    peft_config=peft_config,
    args=training_args,
    train_dataset=load_dataset("json", data_files="watercolor_prompts.jsonl", split="train"),
    reward_funcs=[procedural_watercolor_reward_func]
)

trainer.train()

モード崩壊の回避（KL制御と報酬ハッキング対策）

ビジュアルコードの強化学習では、「全面に単一の濃淡ブルーを塗るだけで高いHPSスコアを獲得できる」といった局所解に陥るモード崩壊（Mode Collapse）が発生しやすい。これを回避するため、以下の2つの制御機構を設ける。

• KLダイバージェンス制御（$\beta$ パラメータ）: GRPOConfigのbetaを0.04前後に設定し、初期方策（SFTモデル）からの過度な乖離を抑止する。学習初期は$\beta=0.02$、ステップ進行に伴い動的に引き上げるコサインスケジュールを適用する。
• マルチモーダル正則化: 単一の美的スコア（HPS）への過剰適合（報酬ハッキング）を防ぐため、VLMによるリファレンス画像とのペアワイズ比較報酬を組み込み、構図の複雑度とテクスチャの多様性を担保する。

ベンチマークと実務でのトレードオフ

8Bクラスモデルをベースとした本パイプラインの計算資源要件と運用上のトレードオフを以下に示す。

本番運用時の落とし穴（Gotchas）

• サンドボックスのセキュリティとリソース枯渇: LLMが意図せず無限ループ（while(true)）や過剰な再帰、メモリを枯渇させる巨大TypedArrayの確保を生成する場合がある。OpenEnvワーカーにはcgroupsによるメモリ上限（512MB）、CPU実行時間クォータ、および2.0秒超過時の厳格なSIGKILL発行とネットワークアクセス完全遮断が必須となる。
• ロールアウト時のVLM APIレートリミット: グループサイズ $G=8$ で並列ロールアウトを行うと、VLM評価APIに対するリクエストが急増し、HTTP 429（Rate Limit）による学習停止を招く。訓練バッチ内の重複描画に対するRedisキャッシュ層の設置、または評価器として軽量なオンプレミスVLM（Qwen2-VL-7B等）への置き換えが実用的である。

導入チェックリストと検証ロードマップ

本アーキテクチャをローカル環境で立ち上げるための段階的検証手順を以下にまとめる。

1. Cold Start用SFTデータの確保: プロンプトと正常に描画されるp5.brushコードのペア（1,000〜2,000件）を準備し、あらかじめSFTを実施しておく（構文エラー率を初期状態で20%以下に抑えることがGRPO収束の必須条件）。
2. 分離サンドボックスコンテナの構築: Node.js + p5.js + p5.brush + headless-gl を内包し、標準入力またはHTTP経由でコードを受け取りPNG画像を返す軽量コンテナを検証する。
3. パイプライン疎通テスト: ランダム報酬を返すダミー関数を用い、GRPOTrainerがOOMを起こさずにステップ更新を完走できるか確認する。
4. HPSスコア単体での最適化確認: r_execとr_aesのみを有効化し、生成されるコードの美的スコアが向上するか、KLダイバージェンスが発散しないかをモニタリングする。
5. VLMペアワイズ評価の統合: リファレンス画像プールとの比較判定を組み込み、特定パターンへのモード崩壊が発生していないか定性検証を行う。

ピクセル生成からコード実行へのパラダイムシフトは、AI画像生成に「編集性」「解像度非依存性」「描画プロセスの透明性」をもたらす。GRPOと分離実行環境を組み合わせた本手法は、リソース制約の厳しい環境下でもプロシージャル生成モデルの自律的品質向上を可能にする実用的なアーキテクチャである。