PaperlayerView exact revision on arXiv

Source-linked AI summary

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

Egor Cherepanov, Daniil Zelezetsky, Alexey K. Kovalev, Aleksandr I. Panov

arXiv:2601.14232v2cs.LGcs.AIcs.CV

TL;DR

Pixel-based RL policies can fail under visual distribution shifts even when latent dynamics and rewards are unchanged, while existing benchmarks entangle shift sources. The paper introduces factorized known-axis environments and benchmark suites, finding strong axis-dependent failures and enabling high-throughput evaluation.

  • Problem

    Existing visual-generalization benchmarks entangle multiple shift sources, making failures difficult to attribute systematically.

  • Method

    KAGE-Env factorizes the observation process into independently controllable visual axes while keeping latent dynamics and rewards fixed, and KAGE-Bench pairs configurations differing along designated axes.

  • Results

    Visual generalization is strongly axis-dependent: filters and effects produce the largest success-rate gaps, while agent appearance is comparatively milder.

  • Takeaways & Limitations

    Known-axis suites enable precise attribution of visual generalization failures under fixed control problems.

  • Takeaways & Limitations

    The benchmark is scoped to a single stylized 2D platformer with fixed dynamics and rewards, and axis rankings may differ for stronger or differently structured agents.

Abstract

from arXiv · show

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/.

1. Introduction

Pixel-based RL policies can fail under purely visual shifts, while existing benchmarks make failure attribution difficult. KAGE-Env and KAGE-Bench isolate visual axes under fixed task structure and support fast systematic evaluation.

  • Motivation: Pixel-based RL policies can degrade sharply under appearance, lighting, and other visual shifts despite unchanged task semantics, dynamics, and rewards.Such brittleness can undermine reliability in robotics, autonomous navigation, and interactive environments.
  • Motivation: Existing benchmarks often entangle visual and structural changes, preventing clean attribution of train–evaluation performance gaps.Simulation cost further limits large-scale ablations and hypothesis testing.
  • Approach: KAGE-Env factorizes observations into independently controllable visual axes while holding latent dynamics and rewards fixed.Performance differences therefore arise from how an observation-based policy responds to different renderings of the same latent states.
  • Scalability: 33M environment steps per second are achieved with KAGE-Env on a single GPU, enabling exhaustive visual-parameter sweeps.The implementation supports 216 parallel environments through JAX compilation and vectorization.
  • Benchmark: KAGE-Bench provides six known-axis suites containing 34 train–evaluation configuration pairs for systematic diagnosis of visual generalization.Each pair targets specified rendering changes while preserving the underlying control problem.

2. Related Work

Prior RL benchmarks establish visual generalization settings but often mix multiple shift sources or limit axis-wise coverage and evaluation scale. KAGE-Bench combines explicit visual factorization with accelerator-native simulation for systematic attribution.

  • Visual generalization in RL: Prior work finds that pixel-based agents often overfit incidental visual features, producing substantial train–test gaps across environments.Spurious correlations include background textures and color statistics rather than task-relevant invariances.
  • Benchmarks: Procgen typically entangles appearance and scene-composition shifts, while DCS exposes explicit distraction axes but covers few factors and is costly to sweep.These design choices limit clean, broad attribution of visual failures.
  • Benchmarks: KAGE-Bench varies specified rendering axes with fixed dynamics and rewards, using a simple platformer to reduce optimization and exploration confounds.Axes include backgrounds, sprites, distractors, filters, and lighting.
  • Fast and scalable evaluation: Visual-generalization evaluation scales combinatorially with N_steps × N_seeds × N_shifts, making CPU-bound simulation expensive.Checkpointing and hyperparameter sweeps compound this cost.
  • Fast and scalable evaluation: KAGE-Env combines accelerator-native execution with observation-process factorization, enabling large-batch, reproducible evaluation of known-axis shifts.This addresses throughput and failure-attribution limitations together.

3. Background

The paper models visual generalization as performance under changed observation kernels with shared latent dynamics and rewards. It supplements return gaps with trajectory metrics because return alone can obscure task competence.

  • POMDP formulation: Each visual configuration defines a POMDP with shared latent states, actions, transitions, rewards, and configuration-dependent observation kernels.Policies receive RGB observations and act without access to the latent simulator state.
  • Visual generalization: The transition kernel and reward function are independent of visual configuration, so visual changes alter observations without changing dynamics or rewards.The same latent state can produce different observations under different configurations.
  • Visual generalization: Visual generalization evaluates a pixel policy trained under D_train and tested under D_eval when only observation kernels change.Both distributions range over visual configurations sharing the same latent control problem.
  • Known-axis shifts: Known-axis evaluation decomposes ξ into a designated visual axis and remaining parameters, attributing differences to that specified observation-process change.Examples include backgrounds, agent sprites, lighting, and filters.
  • Evaluation metrics: Return gaps can obscure competence because they aggregate reward shaping, exploration inefficiency, and penalties, and may remain small when both conditions fail.Trajectory metrics such as distance traveled, normalized progress, and binary success provide finer diagnosis.

4. Known-axis visual generalization

KAGE-Bench formalizes visual generalization by treating rendering changes as induced shifts in a pixel policy’s state-conditional behavior while keeping the latent control problem fixed. This supports axis-isolated evaluation and trajectory-level diagnosis beyond return.

  • Induced state policy: Visual configuration changes affect a reactive pixel policy through the induced state-conditional action distribution obtained by composing rendering and policy kernels.The observation distribution is marginalized before actions are evaluated.
  • Induced state policy: Executing a pixel policy in a visual POMDP has the same state–action process law as executing its induced policy in the latent MDP.The reduction is representational and does not assume policy optimality or modify the control problem.
  • Known-axis attribution: A visual train–evaluation gap equals a performance difference between induced state policies in the same latent MDP.Because transitions and rewards remain unchanged, the gap is attributed to renderer-driven changes in state-conditional behavior.
  • Known-axis attribution: Axis-isolated suites vary one designated visual axis while holding the remaining rendering factors fixed, allowing the measured gap to isolate that axis’s effect.The construction decomposes ξ into the target axis and remaining configuration components.
  • Trajectory-level evaluation: Distance, progress, and success complement return because trajectory-level metrics distinguish partial progress from task completion.Return can mask completion failures when reward shaping provides credit for forward movement.

5. KAGE-Environment

KAGE-Env is a JAX-native 2D platformer whose renderer exposes configurable visual assets and transformations while the latent control interface remains pixel-based. Its configuration and evaluation design support controlled visual variation without changing the underlying task dynamics.

  • Environment design: KAGE-Env parameterizes the renderer with visual configurations while holding the latent control problem fixed.It implements the visual-POMDP interface used for controlled visual-generalization evaluation.
  • Task and interface: The agent receives a single 128×128 RGB image and selects one of eight discrete bitmask actions combining LEFT, RIGHT, and JUMP.Policies interact with the environment exclusively through pixels.
  • Evaluation summary: Table 1 aggregates maximum training-checkpoint metrics across 10 seeds and configurations, reporting distance, progress, success rate, return, and their train–evaluation gaps.Green and red coding indicates smaller and larger gaps, respectively.
  • Rendering assets: Rendering variation includes background images, animated sprite skins, geometric entity shapes, color palettes, and photometric or spatial transformations.The environment provides 128 background images and 27 animated sprite skins, with additional configurable rendering controls.
  • Configuration: All parameters are specified through one YAML configuration organized into background, character, NPC, distractor, filter, effect, layout, and physics groups.The interface exposes both rendering and optional control parameters, while benchmark pairing enforces purely visual isolation.

6. KAGE-Benchmark

KAGE-Bench pairs configurations with identical control problems but controlled rendering changes, organizing 34 train–evaluation pairs into six visual-axis suites. The protocol captures negligible, moderate, and severe transfer gaps using multiple metrics and a standard PPO-CNN diagnostic.

  • Benchmark construction: Each benchmark pair preserves identical transition dynamics and rewards while differing only in a designated subset of rendering parameters.This pairing defines known-axis visual generalization under a fixed control problem.
  • Benchmark construction: A pilot PPO-CNN sweep informed 34 curated train–evaluation pairs spanning agent appearance, background, distractor, effect, filter, and layout suites.Exactly one parameter within the target axis changes in each pair, while easier pairs remain as competence sanity checks.
  • Evaluation protocol: The benchmark uses episodic return and success rate as central metrics, with distance and progress exposing partial or incomplete behavior.Success requires traveled distance to reach 490, near-complete traversal of an approximately 500-distance optimal run.
  • Generalization gaps: KAGE-Bench distinguishes negligible, moderate, and severe generalization gaps, including evaluation collapse despite strong training performance.These regimes describe the relationship between training and evaluation performance across configuration pairs.

7. Results

KAGE-Bench reveals substantial, axis-dependent visual generalization failures under fixed dynamics and rewards. Success-rate degradation is strongest for filters and lighting effects, while several shifts preserve motion yet disrupt task completion.

  • Overall results: 34 train–evaluation configuration pairs across six suites expose persistent train–evaluation gaps under purely visual shifts with fixed dynamics and rewards.Training success rises rapidly, while evaluation success often saturates substantially lower.
  • Axis dependence: Filters cause the largest success-rate degradation (∆SR = 86.8%), followed by effects (80.5%), layout (62.8%), and background (53.3%).Distractors (30.9%) and agent appearance (21.1%) produce comparatively milder gaps.
  • Backgrounds: Background shifts reduce distance by 30.5%, progress by 30.5%, and success rate from 0.90 to 0.42 (∆SR = 53.3%).Evaluation success declines as cumulative background-color support expands beyond the black training background.
  • Photometric and lighting shifts: Filters and effects cause moderate distance degradation (∆Dist = 11.7% and 20.8%) while success rate collapses from 0.83→0.11 and 0.82→0.16.Motion and shaped reward can persist while task success fails under photometric and lighting shifts.
  • Completion versus motion: Distractors and layout show small distance/progress gaps (∼3–4%) but sizable success-rate drops of 30.9% and 62.8%.Increasing same-as-agent distractors from 0 to 11 progressively suppresses evaluation success despite unchanged training success.
  • Per-configuration behavior: Per-configuration failures are heterogeneous, including ∆SR = 98.9% for black→noise backgrounds, 98.8% for hue shift 180◦, and 92.0% for seven same-as-agent distractors.Joint reporting of distance, progress, success rate, and return is needed because both train and evaluation can fail, producing small return gaps.

8. Limitations

KAGE-Bench is a controlled diagnostic benchmark rather than a comprehensive proxy for real-world reinforcement learning. Its conclusions are bounded by a single stylized 2D platformer, fixed dynamics and rewards, and the PPO-CNN probe used to select suites.

  • Scope: KAGE-Bench targets known-axis visual distribution shifts, not comprehensive real-world RL.The benchmark prioritizes clean attribution of visual effects over broad realism.
  • Benchmark setting: The current suite uses one stylized 2D platformer with fixed dynamics and rewards, so reported performance changes are attributable to visual factors within that setting.The visual shifts are factorized rather than photorealistic.
  • Baseline dependence: Axis rankings depend on the representative PPO-CNN probe and may differ for stronger or differently structured agents.The ordering should be interpreted as a property of this controlled benchmark and baseline family, not a universal ranking of visual robustness difficulty.
  • Future scope: Future extensions should add geometric and perspective shifts, richer task families, and complementary benchmarks that trade controlled attribution for broader realism.Double jumps, obstacles, or richer objectives should be treated as new variants because they introduce interacting factors beyond isolated visual shifts.
  • Out of scope: Continual and curriculum learning are natural uses of KAGE-Env but remain outside the present benchmark’s scope.The environment interface supports additional mechanics, but those additions define new benchmark variants.

9. Conclusion

The paper introduces KAGE-Env and KAGE-Bench for fast, controlled evaluation of pixel-based RL under independently configurable visual shifts. Results show that generalization difficulty varies strongly by axis, with background, photometric, and lighting changes producing the most severe completion failures.

  • Conclusion: KAGE-Env factorizes visual observations into independently configurable axes while keeping the underlying control problem fixed and enabling high-throughput simulation.Its JAX-native implementation supports end-to-end compilation and large-scale parallel execution.
  • Conclusion: KAGE-Bench comprises six known-axis suites and 34 train–evaluation configuration pairs that isolate specific sources of visual shift.The design supports precise attribution of performance changes.
  • Conclusion: Background changes and photometric or lighting perturbations induce the most severe failures, often collapsing task success despite nontrivial progress.Agent-appearance shifts are comparatively benign.
  • Conclusion: KAGE-Bench provides a fast, reproducible, diagnostic framework for evaluating pixel-based RL under controlled visual variation.The paper positions it as a basis for systematic robustness analysis and future work on richer shifts and broader task families.

Impact Statement

The paper releases an evaluation environment and configuration suites for visual generalization in reinforcement learning, while noting possible downstream applications of improved autonomous systems.

  • The work introduces an evaluation environment and configuration suites for studying visual generalization in reinforcement learning.
  • KAGE-Bench models each environment instance with a visual configuration ξ that changes rendering from latent states into pixel observations.
  • The induced state policy πξ is obtained by marginalizing the pixel policy over observations rendered from each latent state.
  • Changing visual configuration affects the effective action distribution while leaving the underlying latent control process unchanged.
  • The formal setting comprises latent states, actions, transitions, rewards, observations, visual configurations, and reactive pixel policies.

A.4. Proof of Theorem A.4

The proof shows that marginalizing observations converts a reactive pixel policy into an induced state policy with identical latent state–action behavior and expected return.

  • Proof structure: The visual POMDP generates observations from latent states, samples actions from the reactive pixel policy, and transitions using only the state–action pair.
  • Conditional action law: Conditioning on the latent state and averaging over possible observations yields the induced action distribution πξ(· | st).
  • State–action equivalence: The resulting state–action recursion matches execution of πξ in the latent MDP, so both processes have the same joint law.
  • Interpretation: The equivalence is representational: it does not assert that the induced policy is optimal or that marginalization improves performance.
  • Trajectory-level metrics: Distance, progress, and success metrics inherit the same equivalence because they are measurable functions of the latent trajectory.

C. Additional robustness baseline results

The appendix adds robustness experiments comparing PPO-CNN and RAD-PPO with cutout-color augmentation and illustrates how spatial augmentations can disrupt platformer observations.

  • The appendix compares PPO-CNN with RAD-PPO using cutout-color augmentation on the hardest KAGE-Bench configurations.
  • Random crop-resize augmentation can remove the agent and disrupt spatial structure, hindering distance estimation and jump timing.
  • Random shift disrupts spatial structure at every step, making cliff and obstacle distances and jump timing difficult to estimate.

D. Benchmark Training Details

The benchmark training details report learning curves for all 34 configuration pairs and document the visual parameters used across six generalization axes.

  • Training evaluation: Learning curves cover all 34 train–evaluation configuration pairs, evaluating policies on both in-distribution and out-of-distribution configurations.
  • Metrics and runs: Training metrics include passed distance, progress, success or success-once, and episodic return, with curves summarized across 10 independent runs.
  • Visual axes: The benchmark varies visual parameters for backgrounds, agent sprites and shapes, colors, NPCs, sticky NPCs, and shape distractors.
  • Assets: The configuration specifies a 128×128 screen and includes a background palette of 128 unique scenes.
  • Filters: Photometric settings include brightness, contrast, gamma, saturation, hue shift, and color temperature.
  • Effects: Additional effects include stochastic color jitter and Gaussian noise, pixelation, vignetting, radial light, and point-light controls.
Loading 2601.14232v2…