vlm-gym: RL gym for vision language models (JAX)

TL;DR

vlm-gym is a JAX-based reinforcement learning gym for vision-language models, built around GRPO and a small set of environments. Its headline environment, geospot, casts GeoGuessr-style street-view geolocation as an RL task with a hierarchical reward — country → region → city → coordinates — using geodesic distance with exponential-decay kernels that blend progressively from coarse to fine localization. The reference setup trains Qwen3-VL-4B on the OpenStreetView-5M dataset at lr 5e-7 for 10k steps on an A100-class GPU. A second environment, nlvr2, covers two-image True/False reasoning. The repo is interesting as a clean, self-contained example of applying GRPO to a perception-heavy, continuous-target VLM task using shaped rewards instead of a single scalar.

Key claims

vlm-gym provides a JAX RL gym for VLMs with GRPO training, evaluation, and rollout harnesses, plus an HF→JAX weight converter (utils.hf_to_jax) [README].
The geospot environment scores street-view localization with a four-level hierarchical reward (country, region, city, coords); each level uses geodesic distance passed through an exponential-decay kernel, with weights blending coarse→fine over training [README].
Default training recipe: Qwen3-VL-4B, --env_name geospot, --lr 5e-7, --total_steps 10000, on OpenStreetView-5M [README].
A custom environment is implemented by subclassing envs.base.BaseEnv with reset(idx) → (state, obs) and step(state, action_tokens) → (state, [], reward, done, info) [README].
Hardware target: Linux + CUDA 12, NVIDIA GPU with 80GB+ recommended for training (≈10GB for 4B inference); JAX 0.6.1 [README].
Inference throughput in JAX currently lags the Hugging Face baseline (first-token latency includes XLA compile); optimization is flagged as a TODO [README].
Builds on kvfrans/lmpo (the LM PPO/GRPO trainer it borrows from) and jax-ml/jax-llm-examples for the Qwen3-VL JAX port [README].

Method

The core loop is standard GRPO over VLM rollouts. For geospot, the agent receives a street-view image, emits a text completion that names a country/region/city and/or numeric coordinates, and the environment computes a composite reward by comparing each predicted field against ground truth. Each field’s distance-to-target is mapped to a reward via an exponential-decay kernel (so being “close” on country gets non-trivial credit even when the city is wrong), and a schedule blends the four field weights so early training emphasizes coarse signals (country) and later training emphasizes fine ones (exact coordinates). This is a hand-shaped curriculum baked into the reward rather than the data.

Training and rollout reuse the same Qwen3-VL-4B checkpoint converted from HF to JAX; core/train.py runs GRPO, core/rollout.py runs sampling with standard temperature/top-p/top-k knobs, and core/eval.py benchmarks against the HF baseline. The framework also ships an nlvr2 environment (two-image True/False) to demonstrate the env API on a non-geospatial multimodal-reasoning task.

Results

No paper, no headline metrics. The repo’s only quantitative claim is a “preliminary benchmark” of Qwen3-VL-4B on A100 80GB for inference latency (HF vs. JAX), where the JAX path is currently slower than HF post-compile — this is an engineering snapshot, not a learning curve. The author’s related blog post on a separate online-RL geolocation system (blog.sdan.io/geospot-infinity/) reports that an online RL re-ranker degraded a frozen GeoCLIP baseline by ~414 km / 17% because >65% of users clicked the top guess regardless of accuracy — a useful negative result on the broader “learn from clicks” framing, though it’s not vlm-gym proper.

Why it’s interesting

This is the kind of repo Luma’s RL-for-perception work can crib from directly: it isolates GRPO on a VLM in JAX with a non-trivial, continuous, multi-scale reward (geodesic distance with a coarse-to-fine schedule), which is closer to grounded generative tasks than the discrete-correctness rewards typical of reasoning-RL repos. The hierarchical reward design is the most transferable idea — the same coarse→fine shaping pattern applies to any task with nested target granularity (e.g. layout → pose → pixel for image generation; scene → object → attribute for grounded VQA). It also pairs naturally with Close the Loop: Synthesizing Infinite Tool-Use Data via Multi-Agent Role-Playing, which uses GRPO with gated rewards for tool-use; both repos point at GRPO + shaped rewards as the practical default for non-binary VLM tasks.