KAGE-Bench: Fast Known-Axis Visual Generalization Evaluation for Reinforcement Learning

Abstract: Pixel-based reinforcement learning agents often fail under purely visual distribution shift even when latent dynamics and rewards are unchanged, but existing benchmarks entangle multiple sources of shift and hinder systematic analysis. We introduce KAGE-Env, a JAX-native 2D platformer that factorizes the observation process into independently controllable visual axes while keeping the underlying control problem fixed. By construction, varying a visual axis affects performance only through the induced state-conditional action distribution of a pixel policy, providing a clean abstraction for visual generalization. Building on this environment, we define KAGE-Bench, a benchmark of six known-axis suites comprising 34 train-evaluation configuration pairs that isolate individual visual shifts. Using a standard PPO-CNN baseline, we observe strong axis-dependent failures, with background and photometric shifts often collapsing success, while agent-appearance shifts are comparatively benign. Several shifts preserve forward motion while breaking task completion, showing that return alone can obscure generalization failures. Finally, the fully vectorized JAX implementation enables up to 33M environment steps per second on a single GPU, enabling fast and reproducible sweeps over visual factors. Code: https://avanturist322.github.io/KAGEBench/.

• The paper introduces KAGE-Bench, a diagnostic platform that decouples visual observations from MDP dynamics to isolate axis-specific generalization gaps.
• It demonstrates that controlled visual perturbations can cause drastic performance drops, with background shifts reducing success rates by up to 98.9%.
• The approach scales to millions of environment steps per second, providing a practical tool for robust evaluation and rapid RL algorithm improvements.

Fast Known-Axis Visual Generalization Evaluation with KAGE-Bench

Introduction and Motivation

Visual generalization remains a central unsolved issue in reinforcement learning (RL) from pixels. Despite extensive advances in data augmentation, representation learning, and auxiliary objectives, robust generalization to unseen visual variations—backgrounds, lighting conditions, distractors, and agent appearances—continues to undermine the deployment of RL agents in real-world settings where the observation channel is highly mutable, but task semantics and latent dynamics are invariant. Existing evaluation protocols and benchmarks (e.g., DMC, DCS, Procgen, RL-ViGen) typically entangle multiple factors of variation, which precludes direct attribution of performance degradation to specific visual causes. Moreover, computational overheads in most environments inhibit broad axis-sweep and scalable diagnosis at the necessary granularity.

This paper introduces KAGE-Bench, a JAX-native diagnostic pipeline built atop the KAGE-Env platformer. The platform is explicitly factorized to provide axis-separated observation kernels, enabling precise, high-throughput evaluation of policy sensitivity to controlled visual perturbations, with fixed dynamics and rewards throughout. KAGE-Bench is accompanied by comprehensive suites of 34 paired train/eval configurations, each isolating a visual factor such as backgrounds, photometric filters, agent sprites, distractors, or lighting effects.

Design of KAGE-Env and KAGE-Bench

KAGE-Env is a highly vectorized RL environment parameterized by a set of 93 YAML-specified visual and control parameters. Unlike conventional simulators, the observation kernel $O_\xi(\cdot|s)$ is decoupled from the MDP dynamics. This disjoint parameterization guarantees that latent states ($\mathcal{S}$), actions, transitions ($P$), and reward functions ($r$) remain unchanged regardless of the realized appearance of the environment. The rendering process, including choice of backgrounds, sprite library, photometric and geometric transformations, distractor density, and lighting overlays, is governed exclusively by the configuration $\xi \in \Xi$.

Figure 1: KAGE-Bench factorizes the observation channel into known, independently controllable axes, facilitating axis-isolated generalization diagnosis.

Each suite in KAGE-Bench systematically varies one visual axis per configuration pair, ensuring that any observed generalization gap $$J(\pi;\mathcal{D}_{\mathrm{train}}) - J(\pi;\mathcal{D}_{\mathrm{eval}})$$ reflects only the axis of interest. The environment supports JIT compilation, VMAP evaluation, and can scale to $2^{16}$ parallel environments per accelerator, achieving up to 33 million env-steps/sec.

Theoretical Foundation: Reduction to Induced Policy Shift

The primary theoretical construct is the "induced state policy" $\pi_\xi(a|s) = \int\pi(a|o)O_\xi(do|s)$, which describes the effective state-conditional action distribution resulting from marginalizing the policy over possible observations generated from the same latent state. The central result proves that for any two visual configurations $\xi_1, \xi_2$ (differing only in visual parameters):

$$J(\pi; \mathcal{M}_{\xi_1}) - J(\pi; \mathcal{M}_{\xi_2}) = J(\pi_{\xi_1}; \mathcal{M}) - J(\pi_{\xi_2}; \mathcal{M})$$

All trajectory-based evaluation metrics (distance, normalized progress, binary success) inherit this equivalence by construction. Thus, KAGE-Bench explicitly isolates the perception-induced failure channel.

Figure 2: The induced state policy formalizes the composition of a visual rendering kernel and a reactive pixel policy into a latent-state-conditional action law.

Empirical Results: Axis-Dependent Visual Generalization Gaps

Empirical analyses (PPO-CNN baseline) revealed that different axes induce sharply distinct patterns of policy fragility:

• Background and Photometric Filters: Shifts in background imagery, hue, or color statistics almost universally led to policy collapse, even in settings with trivial underlying latent dynamics. For example, switching from a black to a highly textured or chromatic background reduced success rates by up to 98.9% and resulted in absolute return gaps up to 1800.
• Lighting and Effects: Modulating point light count or strength severely degraded task completion (success rate $\downarrow$ 95.5% in extreme cases), with relatively less impact on overall locomotion metrics (distance, progress), confirming that shaped rewards can mask catastrophic failures.
• Distractors and Layout: Introduction of same-as-agent distractors or layout changes typically resulted in a $30-60\%$ reduction in binary success metrics but only modest progress or distance gaps.
• Agent Appearance: Changes to agent sprites or colors yielded the smallest generalization gaps, with PPO-CNN showing mild reductions in success and negligible drops in progress or distance.

Figure 3: Train-eval performance curves highlighting negligible, moderate, and severe generalization gaps for representative configuration pairs.

In extended axis sweeps, a monotonic dose-response relationship between the difficulty of the shift (e.g., number of distractors, degree of background diversity, intensity of effects) and evaluation success was consistently observed.

Figure 4: Success rate degradation under cumulative single-axis visual shifts: (left) adding background complexity, (right) increasing distractor count produces consistent suppression of evaluation success.

Practical Throughput and Usability

KAGE-Env achieves practical simulation throughput that enables large-scale parameter sweeps and repeated, robust evaluation at the scale previously unattainable for image-based RL. This is critical for rigorous hypothesis testing, hyperparameter tuning, and ablations over visual factors, as required for methods development in robust perception-based RL.

Implications and Future Directions

KAGE-Bench exposes the limiting axis-dependence of visual RL agents: most policies, even with standard augmentations, fail to acquire invariances to visual structure not explicitly modeled by the training distribution. The implication is that current regularization and data augmentation techniques (~e.g., RAD, DrQ) are insufficient to close axis-specific generalization gaps.

From a practical standpoint, the benchmark enables rapid, reproducible diagnosis of generalization defects, essential for iterative algorithmic innovation. Theoretically, the "known-axis" factorization paradigm may inform future benchmarks that combine compositional axes, more structured control problems, and richer observation models.

Future research directions include: leveraging causal inference for robust state abstraction, integrating object-centric and slot attention modules, adversarially composing unseen axes, and exploring cross-benchmark generalization for real-world transfer scenarios.

Conclusion

KAGE-Bench establishes a formal, high-throughput, and axis-disentangled platform for evaluating and diagnosing visual generalization in pixel-based RL. Its design elegantly reduces generalization gaps to policy evaluation in a fixed MDP, enabling unambiguous attribution of observed failures to the perception channel. Empirical analysis demonstrates that visual sensitivity is highly axis-dependent, with the most severe failures found in backgrounds and photometric filters. KAGE-Bench thus provides both a critical foundation and a practical tool for subsequent research into robust, generalizable RL agents.