Training a Conditioned Video Game Agent on a VLM Annotated Dataset

The paper, Training a Conditioned Video Game Agent on a VLM Annotated Dataset by Katrin Schmid and Iuri Frosio (NVIDIA), proposes a method to bypass several practical difficulties of Reinforcement Learning (RL) for video game agents by using Vision Language Models (VLMs) to annotate an offline dataset with human-defined rewards, and then training a conditioned agent via offline RL.

Motivation and Problem. The authors highlight that RL, while powerful, is far from easy-to-use for video games. Key issues include: (1) access to the game engine is required to get rewards for training; (2) rewards are often sparse making RL hard due to the credit assignment problem; and (3) when multiple partial rewards exist, weighting them (e.g., knight vs. pawn in chess) is a difficult trial-and-error process that may not yield explainable policies. Their central insight is that rewards may be difficult to extract from the environment and translate numerically, but easy to describe in words by experts.

Contributions. The paper's contributions are to: (1) demonstrate on a toy problem (Trackmania) that a VLM-annotated offline dataset can be used for offline RL, producing an agent that estimates the same rewards in real time and can be used as cheap annotator; (2) show that a conditioned agent responds to the user requests according to the set of adopted rewards; and (3) study issues that emerged, especially regarding the sparsity of the rewards and the difficulty of training.

Method. Experiments were conducted on Trackmania with a short, visually clean track. A dataset of human gameplay was acquired, consisting of 25,000 frames at 512×512 px resolution, captured at 20 Hz. Frames were annotated with Qwen-3.5 once every 6 frames (0.3 s intervals); the VLM was asked three questions, each answered as a probability in 0–100%:

• Q1: is the red and white car in the middle of the road?

• Q2: is the red and white car moving forward?

• Q3: is the red and white car hitting the lateral barriers?

Each question is associated with a reward Rj. For each reward, a discounted return over a 1.5 s window (30 frames) is defined as r j(t) = Σ k=0 30 R̂ j(t+k) · γ k with γ = 0.99; this is used as conditioning input to the agent, similarly to decision transformers. The model processes the past 6 frames (resampled to 128×128) through three convolution layers with leaky-ReLU activations, and key presses plus desired returns through linear layers; outputs include logits for key presses in the next 2 frames and predictions for future returns/rewards. Training uses binary cross-entropy losses plus a small regularization term for key entropy, RMSProp with learning rate 0.001, batch size 128, and approximately 6000 training steps. Without the return input, the model performs baseline behavioral cloning.

Evaluation and Results. Ground truth rewards were computed using Openplanet, which exposes the car's 3D position and speed. The car's position on the track was used to compute R1 (middle of road) by projecting onto the segment between the closest left/right barrier points, R2 (moving forward) from gear and speed, and R3 (hitting barrier) from distance to the barriers. Results (Table I) show:

• R2 is predicted accurately with F1 > 90%, though it is class-imbalanced (car moving forward 87% of the time).

• R1 is slightly less accurate but clearly learned.

• R3 is problematic: barrier hits occur only 9–11% of the time, making training data unbalanced. Also, a barrier hit may take as little as a few frames, while our VLM annotates the dataset with the probability of the event every 6 frames, so temporal resolution may be insufficient. Additionally, the 128 px resolution may make barrier-hit identification harder (e.g., impact sparkles may not be visible).

The key claim—that conditioning works—is supported by the average rewards: when conditioning to minimize r1, the car is in the middle 37% of the time, versus 43% when maximizing r1; the barrier-hit probability drops from 11% to 9%. Fig. 4 shows the unconditioned agent grabbing curbs on an optimal racing line, while the agent asked to minimize r1 stays closer to track edges, and the maximizer spreads more across the track.

Limitations and Discussion. The authors observe that the conditioning effect tends to vanish in the second part of the track, which they attribute to covariance shift: the model follows instructions where it is easy to stay within the training distribution, but as small errors accumulate, the level of confidence of the model decreases and conditioning no longer works as expected. This is attributed to the simple training algorithm, not the annotation pipeline, and they note that conditioned-model training requires more data than simple behavioral cloning. A promising aspect is that new training data can be automatically generated by a partially trained agent and then annotated by a VLM with no human supervision, analogous to DAGGER but without needing an expert executor—only an expert VLM annotator/supervisor. The authors conclude that, while promising, the approach has limitations: "some rewards may remain sparse and lead to unbalanced training datasets, and attention must be paid to the distribution of returns and training data in order to train conditioned agents that do not go out of distribution at inference time."

Improvements to AI systems

1. Use VLM annotation on short video clips instead of single frames — Query the VLM on small temporal windows (e.g., 5 consecutive frames) to detect transient events like barrier hits, instead of annotating one frame every 6 frames. This fixes the temporal sparsity that made rare rewards (barrier collisions) nearly unlearnable.

2. Add uncertainty-aware reward filtering — Make the VLM output a probability and a confidence estimate. Discard or down-weight low-confidence annotations, and use confidence-weighted losses during offline RL training. This improves reward signal quality and reduces learning from hallucinated or ambiguous labels.

3. Balance sparse-reward training data — Oversample episodes/frames that contain rare events (barrier hits, off-center driving) and apply focal loss or class-balanced reweighting when training reward predictors. This directly addresses the 9–11% barrier-hit imbalance.

4. Use return-conditioned bootstrapping with dropout of conditioning — During training, randomly drop the desired-return conditioning vector (e.g., set it to a null token) so the agent learns to act even when user-specified returns become out-of-distribution. This prevents the observed “conditioning vanishes in the second half of the track.”

5. Add a current-return estimator as policy input — Instead of only receiving desired returns, feed the agent its currently predicted rewards/returns as an observation. This helps the policy detect deviation from user intent and correct itself in real time, reducing covariance shift.

6. Combine offline RL with conservative algorithms (e.g., IQL or CQL) — Replace simple behavioral cloning with an offline-RL algorithm that handles out-of-distribution actions and returns more robustly. This makes the conditioned agent stick to user requests even under distribution shift.

7. Build a self-improving data-generation loop — Deploy a partially trained agent to generate new gameplay data, have the VLM annotate it automatically, append it to the training set, and retrain. This creates a DAGGER-style loop with no human expert needed, only the VLM as supervisor. The agent progressively improves coverage of rare states and edge cases.

8. Precompute denser return targets using interpolation — Instead of a 1.5-second discounted return computed from sparse 0.3-second annotations, linearly interpolate VLM reward probabilities between annotated frames to produce smooth per-frame reward estimates. This gives the agent denser, more stable conditioning signals.

9. Use a multi-task architecture with shared visual features — Train one network to simultaneously predict actions, future rewards, and desired-return satisfaction. The auxiliary reward-prediction heads act as intrinsic regularizers, improving action accuracy and making the reward estimates more reliable at inference time.

10. Condition on natural-language instructions directly — Fine-tune a small text encoder to embed user phrases like “stay in the middle of the road” or “don’t hit barriers” into the return-conditioning vector. The system then translates language to desired rewards automatically.

What the improved AI system can do

• Train from offline human gameplay with zero engine access — No in-game reward functions, no API, no hand-crafted reward weights. The VLM supplies all reward annotations from natural-language questions.

• Obey user requests reliably for the entire episode — It maintains requested behavior (e.g., “stay centered”) even after long time horizons, instead of drifting back to behavioral-cloning behavior.

• Learn rare, sparse events robustly — Barrier hits and other low-frequency events are detected and trained on effectively, because temporal clip annotations and class balancing make them appear consistently in the dataset.

• Quantify its own uncertainty — The system knows when its reward predictions are unreliable and can flag those moments for human review or conservative action.

• Continuously improve itself without human data labeling — It generates new driving data, annotates it with the VLM, retrains, and repeats, gradually covering more tracks and harder scenarios.

• Be adapted to new games or tasks by just changing the VLM questions — Since rewards are described in words, a new task only requires new text prompts and new gameplay data, not new reward engineering.

• Deploy as a cheap real-time annotator — During execution, it predicts all reward probabilities in real time, so it can annotate new datasets for other agents or for human analysis.

• Combine multiple user constraints — It can balance “stay in the middle,” “move forward,” and “avoid barriers” by setting different desired-return values, enabling explainable, user-controlled behavior without manually tuning scalar reward weights.