---
title: 'KAGE-Bench: Visual Generalization in RL'
url: https://www.emergentmind.com/topics/kage-bench
type: topic
---

# KAGE-Bench: Visual Generalization in RL

KAGE-Bench is a high-throughput visual generalization benchmark for pixel-based reinforcement learning (RL) that isolates individual visual factors impacting agent performance. Built atop the KAGE-Env JAX-native 2D platformer, it enables controlled, reproducible evaluation across a set of systematically factorized observation axes, providing precise diagnostics of visual robustness and failure modes under distribution shift [2601.14232].

## 1. Formalization: KAGE-Env and Observation Factorization

KAGE-Env defines a latent Markov Decision Process (MDP) with state space $S$ (positions of platforms, agents, and non-player characters), discrete action space $A \equiv \{0, \ldots, 7\}$ (representing 8-way combinations of Left, Right, and Jump via bitmasking), and a fixed transition kernel $P(s'|s,a)$ implementing 2D platformer physics. The reward function
\[
r_t = \alpha_1 \cdot \max(0, x_{t+1} - x_t^{\max}) - \left[ \alpha_2\,\mathbb{I}_{\{\text{Jump}(a_t)\}} + \alpha_3 + \alpha_4\,\mathbb{I}_{\{\text{idling}(x_t, x_{t+1})\}} \right]
\]
shapes for forward progress, penalizes jumping and idling, and enforces an episode length $T=500$ with only timeouts for termination.

Observation rendering is parameterized by a configuration vector $\xi = (A_1, ..., A_K)$, each element $A_k$ controlling a specific visual axis (e.g., background, agent sprites, photometric filters). The rendered pixel observation $o \in \{0,\ldots,255\}^{H\times W \times 3}$ is produced by a deterministic or randomized function $o = f(s; A_1, ..., A_K)$. The observation kernel is $O_\xi(o|s) = \delta(o - f(s; \xi))$, ensuring that all axis controls affect only visual appearance, without altering latent dynamics or rewards.

A pixel policy $\pi(a|o)$ induces a state-conditional policy under visual parameterization $\xi$:
\[
\pi_\xi(a|s) = \int_{\mathcal{O}} \pi(a|o)\,O_\xi(do|s)
\]
and hence the joint law on $(s_t, a_t)$ is preserved under the shift from $M_\xi = (S, A, P, r, O_\xi, ...)$ to $M = (S, A, P, r)$ by running $\pi_\xi$ in the latter. Therefore, any difference in expected return $J(\pi; M_\xi) - J(\pi; M_{\xi'})$ quantifies the policy’s axis-specific visual generalization.

## 2. Known-Axis Benchmark Suite Construction

KAGE-Bench specifies six families of visual axes, each parameterizing distinct, independently controllable components of the rendering pipeline:
1. **Agent Appearance**: geometric shape (circle, line), color, or sprite skin.
2. **Background**: solid colors, random noise, curated images.
3. **Distractors**: animated or static extraneous objects (e.g., NPC skeletons, agent-shaped imposters).
4. **Lighting/Effects**: dynamic effects such as global/point lights, varying intensity or falloff parameters.
5. **Filters**: photometric operations (brightness, contrast, gamma, hue, Gaussian noise, pixelation, vignetting).
6. **Layout**: platform color palette changes (e.g., cyan palette to red palette).

For each axis, the benchmark enumerates a set of train–eval configuration pairs $(\xi^{\text{train}}, \xi^{\text{eval}})$, ensuring that only a single axis changes (all other rendering/dynamics parameters are fixed). In total, 34 such pairs are defined, enabling axis-resolved measurement of generalization gaps.

## 3. Experimental Protocol and Evaluation Metrics

The canonical agent is a PPO-CNN baseline from CleanRL, deployed with 128 parallel environments and batch size $16{,}384$, employing a 3-layer convolutional encoder and fully connected policy/value heads. Each configuration pair is evaluated over 10 random seeds, recording multiple trajectory-level and return-based metrics:
- **Maximum Return ($J$)**: peak mean episodic reward over all training checkpoints.
- **Passed Distance ($F_{\text{dist}}$)**: $x_T - x_0$, absolute horizontal progress.
- **Progress ($F_{\text{prog}}$)**: normalized as fraction of required course length.
- **Success Rate ($F_{\text{succ}}$)**: proportion of episodes completing the task ($F_{\text{dist}} \geq D$).

Generalization gap per metric is defined as the absolute or relative difference between train and evaluation performance (e.g., $\Delta\text{SR} = \text{SR}_\text{train} - \text{SR}_\text{eval}$ and $\Delta\text{Dist}\% = 100{\cdot}(\text{Dist}_\text{train}-\text{Dist}_\text{eval})/\text{Dist}_\text{train}$). The “maximum-over-training” protocol avoids checkpoint-selection artifacts by using the best checkpoint for each seed.

## 4. Analysis of Visual Generalization and Failure Modes

Axis-level aggregation yields the following core findings:
- **Filters and Effects**: Despite only moderate drops in forward motion ($\Delta\text{Dist}\approx12$–$21\%$), success rates collapse almost entirely ($\Delta\text{SR}\approx80$–$87\%$). Policies under photometric and lighting shifts often move forward but fail to complete the course, indicating that measuring only return or distance can mask severe failures.
- **Background Shifts**: Large degradations in both progress (30.5%) and success rate (53.3%). For example, switching from a black to a noise background induces a $\Delta\text{SR}=98.9\%$ drop.
- **Distractors and Layout**: Minor distance gaps ($<$5%) but major drops in completion rate ($\Delta\text{SR}=31\%,\,63\%$). Increasing “same-as-agent” distractors incrementally degrades SR while maintaining near-normal distances.
- **Agent Appearance**: The least harmful axis; shape or color changes yield modest SR reductions ($21.1\%$), confirming such perturbations are comparatively benign.

The benchmark exposes instances where return is decoupled from task completion, especially under strong photometric/lighting shift: agents traverse substantial fractions of the course but rarely succeed according to the defined criteria.

## 5. Vectorized JAX Implementation and Throughput

KAGE-Env is fully implemented in JAX (including both simulator dynamics and rendering). All randomness is managed via explicit PRNG keys, with environment interaction (reset, step) vectorized through `jax.vmap`, and multi-step rollouts via `jax.lax.scan`. This design enables single-call batched simulation of $2^{16} = 65536$ parallel environments on one modern GPU.

Performance measurements report sustained peak throughput up to $33\times10^6$ steps/sec on NVIDIA H100 hardware, with all rendering axes combined. Configurations of comparable complexity consistently yield $10$–$30$ million steps/sec, enabling large-scale ablation studies and multi-seed sweeps in minutes.

## 6. Applications, Limitations, and Future Directions

KAGE-Bench provides a uniquely diagnostic platform for:
- **Quantitative visual generalization analysis**: Its axis-isolated pairs expose policy sensitivities that are otherwise conflated in existing benchmarks.
- **Ablation and robustness evaluation**: The high simulation throughput supports systematic method comparisons and perturbation studies at scale.
- **Method development**: Facilitates fair, reproducible testing of augmentation, invariance, and contrastive learning techniques within a factorized visual space.

Limitations:
- The current suite uses a single 2D platforming scenario; tasks such as navigation or object manipulation remain unexplored.
- Selection of axes is hand-designed; auto-discovery or continuous axis interpolation is not presently addressed.
- Integrating state-of-the-art robust RL methods into the KAGE-Bench framework is ongoing.

A plausible implication is that KAGE-Bench's precise factorization of visual confounds and its scalable implementation may become a reference diagnostic standard for future pixel-based RL research [2601.14232].

Source: https://www.emergentmind.com/topics/kage-bench