1 Introduction↩︎

Reinforcement learning (RL) agents trained from high-dimensional pixel observations are brittle to changes in appearance, lighting, and other visual nuisance factors [1]–[3]. Policies that perform well in-distribution can degrade sharply under purely visual distribution shifts, even when task semantics, transition dynamics, and rewards are unchanged [4]–[6]. This brittleness poses a fundamental obstacle to real-world deployment, where observations inevitably vary due to viewpoint changes, illumination, surface appearance, and sensor noise while the control-relevant latent state remains fixed [7]–[12]. As a result, pixel-based RL policies that rely on incidental visual correlations can fail abruptly despite convergence, undermining reliability in robotics, autonomous navigation, and interactive environments [2], [13]. More broadly, visual generalization is needed wherever models must robustly extract information from visual structure, even in scientific texts and figures [14].

Despite substantial progress in representation learning [15]–[17] and data augmentation [7], [18], [19], understanding visual generalization failures remains challenging. A central obstacle lies in evaluation benchmarks, which often entangle multiple visual and structural changes such as background appearance, geometry, dynamics, and distractors [2], [13], [20]. In these settings, train–evaluation performance gaps cannot be cleanly attributed to specific sources of shift, and failures may reflect visual sensitivity, altered task structure, or interactions between confounded factors. Compounding this issue, many pixel-based RL environments are computationally expensive to simulate, limiting large-scale ablations and slowing hypothesis testing.

We address these limitations with KAGE-Bench (Known-Axis Generalization Evaluation Benchmark), a visual generalization benchmark in which sources of distribution shift are isolated by construction. KAGE-Bench is built on KAGE-Env (1), a JAX-native [21] 2D platformer whose observation process is factorized into independently controllable visual axes while latent dynamics and rewards are held fixed (see 2). Under this known-axis design, each axis corresponds to a well-defined component of the observation kernel, and any train–evaluation performance difference arises solely from how a fixed observation-based policy responds to different renderings of the same latent states, enabling unambiguous attribution of visual generalization failures.

Systematic analysis of visual generalization requires evaluating many controlled shifts at scale. KAGE-Env is implemented entirely in JAX with end-to-end jit compilation and vectorized execution via vmap and lax.scan, enabling efficient large-batch simulation on a single accelerator. In practice, this design scales up to \(2^{16}\) parallel environments on one GPU and achieves up to 33M environment steps per second (see 3), making exhaustive sweeps over visual parameters and fine-grained diagnosis of generalization behavior feasible.

Building on this environment, we construct six visual generalization suites comprising 34 train–evaluation configuration pairs, each targeting a specific visual axis. Using these suites, we demonstrate that visual generalization is strongly axis-dependent and identify classes of visual shifts that reliably induce severe performance degradation, even for a standard PPO-CNN baseline [22]. Representative negligible, moderate, and severe train–evaluation gaps are previewed in 4.

2 Related Work↩︎

2.0.0.1 Visual generalization in RL.

Visual generalization studies whether policies trained from pixel observations retain performance when the observation process changes while latent dynamics and rewards remain fixed. Prior work shows that agents often overfit incidental visual features, leading to substantial train–test gaps across a wide range of environments and settings [17], [23]–[25]. A common explanation is that standard architectures and objectives exploit spurious visual correlations, such as background textures or color statistics, rather than learning task-relevant invariances [13], [19], [20]. Accordingly, many approaches have been proposed to improve robustness, including data augmentation, auxiliary representation learning objectives, and regularization methods [7], [15], [16], [18], [26]–[32]. KAGE-Env and KAGE-Bench provide diagnostic infrastructure for this literature by enabling fast, controlled, axis-specific evaluation that isolates changes in the observation kernel.

2.0.0.2 Benchmarks for visual generalization in RL.

A range of benchmarks study visual generalization in pixel-based RL, differing in task domains and in how explicitly they isolate sources of visual variation. RL-ViGen [2] spans multiple domains, including locomotion, manipulation, navigation, and driving, with shifts in textures, lighting, viewpoints, layouts, and embodiments. [19] evaluates continuous control under controlled appearance changes such as color randomization and dynamic video backgrounds. Obstacle Tower [33] and LevDoom [34] consider 3D settings where many factors vary jointly, making attribution of failures to specific visual causes difficult. Related benchmarks such as DMC-VB [17] and Distracting MetaWorld [35] introduce task-irrelevant visual distractors while keeping task dynamics fixed.

Among widely used benchmarks, Procgen [20] relies on procedural generation, so train–test gaps typically reflect entangled shifts in appearance and scene composition rather than isolated visual factors. The Distracting Control Suite (DCS) [13] introduces explicit distraction axes but is limited to a small set of factors, and broad axis-wise sweeps are costly in its underlying continuous-control simulator. KAGE-Env and KAGE-Bench complement these benchmarks by explicitly factorizing the observation process into independently controllable visual axes. KAGE-Env uses a simple platformer to reduce optimization and exploration confounds, while KAGE-Bench constructs train–evaluation splits that vary specified axes (e.g., backgrounds, sprites, distractors, filters, and lighting) with fixed dynamics and rewards, enabling systematic, axis-specific attribution of generalization failures. Compared with DMControl Generalization and DCS-style evaluation, KAGE-Bench emphasizes a fast JAX-native pixel environment where axis-isolated visual shifts can be swept systematically rather than sampled from a small number of expensive simulator variants. This also addresses a practical gap in the JAX RL ecosystem. MuJoCo Playground [36] provides MJX-based robot-learning environments with state and pixel inputs, and Octax [37] provides high-throughput JAX arcade environments. KAGE-Env is complementary because it targets controllable RGB rendering axes and paired train–evaluation shifts for known-axis attribution. Combining MJX-style simulators with similarly controlled visual-shift layers remains a promising future direction.

2.0.0.3 Fast and scalable evaluation in RL.

Evaluating generalization in RL is sample intensive, as reliable conclusions require averaging over random seeds, environment instances, and distribution shifts. In visual generalization benchmarks, this leads to combinatorial scaling \(N_{\text{steps}} \times N_{\text{seeds}} \times N_{\text{shifts}}\), often compounded by checkpointing and hyperparameter sweeps, making evaluation costly in CPU-bound simulators.

Recent work addresses this bottleneck through accelerator-native RL systems, where environment stepping is implemented as compiled, vectorized computation on GPUs or TPUs. Examples include JAX-based simulators such as Brax [38], Jumanji [39], XLand-MiniGrid [40], CAMAR [41], and Craftax [42], as well as GPU-native platforms such as ManiSkill3 [43], MIKASA-Robo [44], and WarpDrive [45]. By eliminating host-side control flow, these systems achieve orders-of-magnitude throughput. However, high throughput alone does not yield diagnostic evaluation of visual robustness. Benchmarks such as Procgen and DCS do not support exhaustive, axis-isolated sweeps over rendering factors, limiting failure attribution. KAGE-Env combines the accelerator-native paradigm with explicit factorization of the observation process into independently controllable axes, enabling large-batch, reproducible evaluation of known-axis visual shifts under fixed latent dynamics and rewards. This distinction is practical: raw KAGE-Env stepping reaches 33M environment steps per second, while the full PPO-CNN training pipeline used for the benchmark runs at about 52k training steps per second per GPU after CNN forward passes, backpropagation, optimizer updates, logging, checkpointing, and video generation on \(128\times128\times3\) RGB inputs.

3 Background↩︎

3.0.0.1 Partially Observable Markov Decision Processes.

We consider episodic control with horizon \(T\) in a partially observable Markov decision process (POMDP). Each environment instance is indexed by a visual configuration \(\xi \in \Xi\) and defined as \(\mathcal{M}_\xi = (\mathcal{S}, \mathcal{A}, P, r, \Omega, O_\xi, \rho_0, \gamma)\), where \(\mathcal{S}\) is the latent (control-relevant) state space, \(\mathcal{A}\) is the action space, \(P(\cdot \mid s,a)\) is the transition kernel, \(r(s,a)\) is the reward function, \(\Omega\) is the observation space, \(O_\xi(\cdot \mid s)\) is the observation (rendering) kernel parameterized by \(\xi\), \(\rho_0\) is the initial state distribution, and \(\gamma \in [0,1)\) is the discount factor. At each timestep \(t\), the environment occupies a latent state \(s_t \in \mathcal{S}\). An observation is generated according to \(o_t \sim O_\xi(\cdot \mid s_t),~ o_t \in \Omega \subseteq \{0,\dots,255\}^{H \times W \times 3}.\) Based on this observation, the agent selects an action \(a_t \in \mathcal{A}\), receives reward \(r(s_t,a_t)\), and transitions to \(s_{t+1} \sim P(\cdot \mid s_t,a_t)\).

A key structural property enforced throughout this work is that the transition kernel \(P\) and reward function \(r\) are independent of the visual configuration \(\xi\). All dependence on \(\xi\) is confined to the observation kernel \(O_\xi\). Consequently, the same latent state \(s_t\) may give rise to different observations under different values of \(\xi\), while inducing identical dynamics and rewards. Visual generalization concerns the behavior of policies under such changes in the observation process, with the underlying control problem held fixed.

3.0.0.2 Policies and return.

We focus on reactive pixel-based policies that map observations directly to action distributions: \(\pi(a \mid o)\). The expected discounted return of a policy \(\pi\) in environment \(\mathcal{M}_\xi\) is \[J(\pi;\mathcal{M}_\xi) \;=\; \mathbb{E}_{\substack{s_0 \sim \rho_0, \\ o_t \sim O_\xi(\cdot \mid s_t), \\ a_t \sim \pi(\cdot \mid o_t), \\ P}}\!\left[\sum_{t=0}^{T-1} \gamma^t r(s_t,a_t)\right].\]

3.0.0.3 Visual generalization.

We study generalization under shifts in visual parameters that affect observations but not the underlying control problem. Let \(\Xi\) denote the space of visual configurations, and let \(\mathcal{D}_{\mathrm{train}}\) and \(\mathcal{D}_{\mathrm{eval}}\) be probability distributions over \(\Xi\). Each \(\xi \in \Xi\) induces a visual POMDP \(\mathcal{M}_\xi\) through its observation kernel \(O_\xi\), while sharing the same latent dynamics \(P\) and reward function \(r\).

A pixel policy \(\pi(a\mid o)\) is trained using environments with \(\xi \sim \mathcal{D}_{\mathrm{train}}\) and evaluated under \(\xi \sim \mathcal{D}_{\mathrm{eval}}\). For any distribution \(\mathcal{D}\) over \(\Xi\), we define the expected performance \[J(\pi;\mathcal{D}) \;=\; \mathbb{E}_{\xi \sim \mathcal{D}} \!\left[ J(\pi;\mathcal{M}_\xi) \right].\]

We refer to this setting as visual generalization when the shift from \(\mathcal{D}_{\mathrm{train}}\) to \(\mathcal{D}_{\mathrm{eval}}\) changes only the observation kernels \(O_\xi\), while preserving the latent state space, transition dynamics, and reward function.

3.0.0.4 Known-axis visual shifts.

KAGE-Bench focuses on known-axis visual generalization. Each visual configuration is decomposed as \(\xi = (\xi_{\mathrm{axis}}, \xi_{\mathrm{rest}})\), where \(\xi_{\mathrm{axis}}\) specifies a designated axis of visual variation (e.g., background appearance, agent sprites, lighting, filters), and \(\xi_{\mathrm{rest}}\) contains all remaining parameters. By construction, any performance difference between training and evaluation can therefore be attributed to changes in the observation process along the specified visual axis, rather than to changes in task structure, dynamics, or rewards. This intuition is formalized and justified in 4 and 10.

3.0.0.5 Evaluation metrics.

Given \(\mathcal{D}_{\mathrm{train}}\) and \(\mathcal{D}_{\mathrm{eval}}\), we report in-distribution and out-of-distribution performance, \(J(\pi;\mathcal{D}_{\mathrm{train}})\) and \(J(\pi;\mathcal{D}_{\mathrm{eval}})\), and define the return-based generalization gap \[\Delta(\pi) \;=\; J(\pi;\mathcal{D}_{\mathrm{train}}) - J(\pi;\mathcal{D}_{\mathrm{eval}}).\] While \(\Delta(\pi)\) provides a coarse measure of performance degradation under visual shift, it is insufficient to fully characterize generalization behavior. The discounted return aggregates multiple effects, including reward shaping, exploration inefficiency, and penalty terms, and may obscure whether an agent nearly solves the task or fails catastrophically. In particular, if a policy fails under both training and evaluation configurations, the return gap can be small despite the absence of task competence.

For this reason, we complement return-based evaluation with additional trajectory-level metrics that are measurable functions of the latent state trajectory, including distance traveled, normalized progress toward the goal, and binary task success. These metrics distinguish partial progress from complete failure and provide a more fine-grained view of visual generalization behavior. Their precise definitions and empirical use are described in 6.

4 Known-axis visual generalization↩︎

This section states the formal principle behind KAGE-Bench. In our construction (3), the latent control problem is fixed and only the renderer changes: \(\xi\) affects performance only through the induced state-conditional action law obtained by composing the observation kernel with the pixel policy. The goal is to make this channel explicit and to justify the benchmark protocol: (i) constructing suites that intervene on a single visual axis, and (ii) evaluating not only return but also trajectory-level metrics such as distance, progress, and success.

4.0.0.1 From pixel policies to state-conditional behavior.

A reactive pixel policy \(\pi(\cdot\mid o)\) maps observations to actions and does not directly specify an action distribution conditioned on the latent state \(s\). However, in a visual POMDP \(\mathcal{M}_\xi\), the observation kernel \(O_\xi(\cdot\mid s)\) induces a distribution over rendered observations for each latent state. Composing these kernels yields a well-defined state-conditional action distribution by marginalizing the intermediate observation: \[s \xrightarrow{~O_\xi(\cdot\mid s)~} o \xrightarrow{~\pi(\cdot\mid o)~} a.\] Under our construction (and for reactive policies), this composition is the only mechanism by which the visual configuration \(\xi\) can affect control, since \(P\) and \(r\) are invariant across \(\xi\). 5 illustrates this marginalization in a concrete discrete example.

For a fixed pixel policy \(\pi\), the map \(\xi\mapsto \pi_\xi\) summarizes the effect of visual variation on state-conditional behavior. In particular, changing \(\xi\) changes \(\pi_\xi\) while leaving the latent control problem \((\mathcal{S},\mathcal{A},P,r,\rho_0,\gamma)\) unchanged.

4.0.0.2 Visual shift is equivalent to induced policy shift.

The next theorem formalizes the reduction used throughout KAGE-Bench: executing \(\pi\) in the visual POMDP \(\mathcal{M}_\xi\) induces the same latent state–action law as executing \(\pi_\xi\) in the latent MDP \(\mathcal{M}\).

1 is purely representational: it does not assume optimality and it does not modify the control problem. A useful consequence is the identity, for any \(\xi,\xi'\in\Xi\), \[J(\pi;\mathcal{M}_\xi)-J(\pi;\mathcal{M}_{\xi'}) \;=\; J(\pi_\xi;\mathcal{M})-J(\pi_{\xi'};\mathcal{M}), \label{eq:gap95identity}\tag{1}\] which states that a visual train–evaluation gap for a fixed pixel policy \(\pi\) is exactly a performance difference between induced state policies in the same latent MDP. This is the formal basis for attributing failures to the observation process: since \((P,r)\) are unchanged, any degradation under \(\xi\to\xi'\) must be explained by how the renderer changes the induced state-conditional behavior \(\pi_\xi\).

4.0.0.3 Why known-axis suites enable axis-specific attribution.

KAGE-Bench constructs axis-isolated suites by decomposing \(\xi=(\xi_{\mathrm{axis}},\xi_{\mathrm{rest}})\) and pairing train and evaluation configurations that differ only in the designated axis: \[\xi^{\mathrm{train}}=(\xi^{\mathrm{train}}_{\mathrm{axis}},\xi_{\mathrm{rest}}) \;\;\text{and} \;\; \xi^{\mathrm{eval}}=(\xi^{\mathrm{eval}}_{\mathrm{axis}},\xi_{\mathrm{rest}}).\] Equivalently, for all \(s\in S\) the paired renderers satisfy \(O_{\xi^{\mathrm{train}}}(\cdot\mid s)=O(\cdot\mid s;\xi^{\mathrm{train}}_{\mathrm{axis}},\xi_{\mathrm{rest}})\) and \(O_{\xi^{\mathrm{eval}}}(\cdot\mid s)=O(\cdot\mid s;\xi^{\mathrm{eval}}_{\mathrm{axis}},\xi_{\mathrm{rest}})\), so the only change in the observation process is along \(\xi_{\mathrm{axis}}\). Under this controlled-intervention design, the induced policies \(\pi_{\xi^{\mathrm{train}}}\) and \(\pi_{\xi^{\mathrm{eval}}}\) differ only through this axis-dependent change in \(O_\xi\). Therefore, by 1 , the measured gap isolates how that visual axis perturbs the induced state-conditional behavior of \(\pi\).

4.0.0.4 Trajectory-level consequences and evaluation metrics.

By Item 2 of 1, the latent state–action trajectory has the same law under \((\mathcal{M}_\xi,\pi)\) and \((\mathcal{M},\pi_\xi)\), so the reduction applies to any measurable trajectory functional, not only return. We therefore report distance, progress, and success in addition to episodic return: these are functions of the latent trajectory exposed by KAGE-Env for evaluation, and their gaps under \(\xi\!\to\!\xi'\) admit the same induced-policy interpretation. Unlike return, which can mask completion failures due to reward shaping, these metrics separate partial progress from task completion.

All proofs are deferred to 10.

import jax
from kage_bench import (
  KAGE_Env,
  load_config_from_yaml,
)
# Create environment with custom config
env = KAGE_Env(
  load_config_from_yaml("custom_config.yaml")
)
# Vectorize and JIT compile
reset_vec = jax.jit(jax.vmap(env.reset))
step_vec = jax.jit(jax.vmap(env.step))
# Initialize 65,536 parallel environments
N_ENVS = 2**16
keys = jax.random.split(
  jax.random.PRNGKey(42), N_ENVS
)
# Reset all at once
obs, info = reset_vec(keys)
states = info["state"]
# Parallel step: Samples one random discrete
# action per env in [0, 7] (bitmask actions)
actions = jax.random.randint(
  keys[0], (N_ENVS,), 0, 8
)
# obs.shape: (65536, 128, 128, 3)
obs, rewards, terms, truncs, info \;
  = step_vec(states, actions)
states = info["state"]

5 KAGE-Environment↩︎

KAGE-Env (1, [code:env]) is a JAX-native RL environment designed for controlled evaluation of visual generalization. It implements the visual-POMDP interface from 3: configurations \(\xi\in\Xi\) parameterize the renderer \(O_\xi(\cdot\mid s)\) while the latent control problem is held fixed.

5.0.0.1 Task and interface.

KAGE-Env is an episodic 2D side-scrolling platformer with horizon \(T\) and a push-scrolling camera. At each timestep \(t\), the agent observes a single RGB image \(o_t \in \{0,\dots,255\}^{H\times W\times 3}\), with default resolution \(H=W=128\), and selects an action \(a_t\) from a discrete action space \(\mathcal{A}=\{0,\dots,7\}\). Actions are encoded as a bitmask over three primitives: \(\mathrm{\small Left}=1\), \(\mathrm{\small Right}=2\), and \(\mathrm{\small Jump}=4\). Policies interact with the environment exclusively through pixels; the latent simulator state is not available to the policy and is exposed only via the info dictionary for logging and evaluation. The Left action is retained to preserve a conventional platformer-style action space, although a competent policy can learn to ignore it when moving left is not useful.

background:
  mode: "image"
  image_paths:
    - "src/kage/assets/backgrounds/bg-1.jpeg"
    - "src/kage/assets/backgrounds/bg-64.jpeg"
    - "src/kage/assets/backgrounds/bg-128.jpeg"
  parallax_factor: 0.5
  switch_frequency: 0.0
character:
  mode: "sprite"
  sprite_paths:
    - "src/kage/assets/sprites/clown"
    - "src/kage/assets/sprites/skeleton"
  enable_animation: true
  animation_fps: 12.0
npc:
  mode: "sprite"
  sprite_dir: "src/kage/assets/sprites"
filters:
  brightness: 0.0
  hue_shift: 0.0

5.0.0.2 Reward and termination.

Let \(x_t\in\mathbb{R}\) denote horizontal position and \(x_t^{\max} := \max_{0\le k\le t} x_k\) the furthest position reached so far. The per-step reward is \[\label{eq:kage95reward}\begin{align} &r_t=\underbrace{\alpha_1 \max\!\left\{0,\, x_{t+1}-x^{\max}_t\right\}}_{\text{first-time forward progress}} \\ - &\underbrace{\Big(\alpha_2\,\mathbb{I}[\mathrm{\small Jump}(a_t)] + \alpha_3 + \alpha_4\,\mathbb{I}[\mathrm{idle}(x_t,x_{t+1})]\Big)}_{\text{penalties}} , \end{align}\qquad{(1)}\] where \(\mathbb{I}[\mathrm{\small Jump}(a_t)]\) indicates the jump bit is active in \(a_t\), \(\alpha_3\) is a per-timestep time cost, and \(\mathrm{idle}(x_t,x_{t+1})\) flags lack of horizontal progress. Episodes terminate only by time-limit truncation at \(T=\texttt{episode\_length}\).

5.0.0.3 Rendering assets and visual parameters.

KAGE-Env provides a library of visual assets and rendering controls for constructing visual variation. Assets include 128 background images (Appendix, 46) and 27 animated sprite skins for the agent and non-player characters (Appendix, 47); when sprites are disabled, entities can be rendered as geometric shapes (9 types) with a palette of 21 colors. The renderer further exposes photometric and spatial transformations (e.g., brightness, contrast, gamma, hue, blur, noise, pixelation, vignetting) and lighting/overlay effects such as dynamic point lights with configurable count, intensity, radius, falloff, and color.

5.0.0.4 Configuration interface.

All parameters are specified through a single .yaml configuration file ([code:config]). A configuration \(\xi\in\Xi\) is organized into groups background, character, npc, distractors, filters, effects, layout, and physics. These groups include rendering parameters (affecting only \(O_\xi\)) as well as optional control parameters (affecting \(P\) or \(r\)). KAGE-Env exposes both for extensibility; isolation of purely visual shifts is enforced by the KAGE-Bench pairing protocol (6).

6 KAGE-Benchmark↩︎

KAGE-Bench is a benchmark protocol built on top of KAGE-Env. It specifies how environment configurations are selected and paired to evaluate known-axis visual generalization. Concretely, KAGE-Bench defines a set of train–evaluation configuration pairs \((\xi^{\mathrm{train}},\xi^{\mathrm{eval}})\) such that the underlying control problem is identical (\(P^{\mathrm{train}}=P^{\mathrm{eval}}\), \(r^{\mathrm{train}}=r^{\mathrm{eval}}\)) and the two configurations differ only in a designated subset of rendering parameters.

6.0.0.1 Benchmark construction.

We first conduct a pilot sweep over KAGE-Env’s rendering parameters using a standard PPO-CNN, adopted from the CleanRL [46] library¹, trained from a single RGB frame without history or frame stacking. Hyperparameters are reported in Appendix, ¿tbl:tab:ppo95cnn95hparams?. This sweep measures how individual rendering parameters affect out-of-distribution performance when the control problem is fixed. Based on these results, we curate 34 train–evaluation configuration pairs that exhibit a range of generalization behavior, including both severe and mild gaps. The selected pairs are grouped into six suites corresponding to distinct visual axes: agent appearance, background, distractors, effects, filters, and layout. In each pair, exactly one parameter within the target axis is changed between train and evaluation, while all other parameters are held fixed. Easier pairs are intentionally retained as sanity checks, ensuring that the benchmark distinguishes lack of generalization from lack of task competence. The 34 pairs should be interpreted as a clear reference subset, not as a fixed universal difficulty ranking. PPO-CNN is used as a diagnostic probe to expose which controlled visual shifts break a standard pixel policy; different methods may expose different failure modes, and KAGE-Env configurations can be extended as methods improve. For additional robustness validation, we also evaluated baselines inspired by common visual-RL robustness methods: RAD-PPO with cutout-color augmentation, RAD-PPO with random crop-resize augmentation [18], CURL [47], DrQ [8], and SVEA [48]. These baselines used official PyTorch implementations interfaced with KAGE-Env through a JAX-to-PyTorch wrapper and were evaluated on the most challenging configurations identified in the PPO-CNN sweep.

6.0.0.2 Evaluation protocol and metrics.

For each train–evaluation configuration pair, we run 10 independent training seeds and periodically evaluate the current policy on both configurations. For each run and metric, we record the maximum value attained over training, average these maxima across seeds to obtain per-configuration results, and then average within each suite to produce the axis-level summaries in 1. We use the maximum-over-training statistic to assess whether a visual generalization gap is in principle mitigable by a given method. Because generalization performance can be non-monotonic and peak at different iterations across runs, this aggregation provides an upper envelope on achievable transfer and avoids confounding results with arbitrary checkpoint selection. Each seed is trained for 25M environment steps and evaluated every 300 PPO iterations on both the training and evaluation configurations using 128 evaluation episodes per checkpoint. We use episodic return and success rate as central metrics, while distance and progress provide interpretability for partial progress and incomplete behavior. Progress normalizes distance by dist_to_success; success is triggered when the traveled distance reaches dist_to_success \(=490\), which corresponds to near-complete traversal of an approximately 500-distance optimal run and was not tuned to exaggerate gaps. The full benchmark comprises \(34\) configuration pairs \(\times\) \(10\) seeds, requiring about 500 GPU-hours on four A100 GPUs, or roughly 1.5 hours per run including checkpointing and rollout videos savings. Evaluating a new method does not require all 34 pairs: the hard subset used for the RAD/CURL/DrQ/SVEA experiments is sufficient for a cheaper stress test.

6.0.0.3 Generalization gap.

We define the visual generalization gap as the performance difference between the training and evaluation configurations of a pair. 4 illustrates three characteristic regimes observed in KAGE-Bench: (i) negligible gap, where train and eval performance coincide; (ii) moderate gap, where partial transfer occurs; and (iii) severe gap, where evaluation performance collapses despite strong training performance. Full learning curves for all 34 configuration pairs and all suites are reported in 13.

2.5pt 0pt 0pt

7 Results↩︎

1 reports axis-level results for PPO-CNN under our maximum-over-training protocol: for each seed we take the maximum of each metric over training checkpoints, then average across 10 seeds and finally across configuration pairs within an axis. 6 complements this summary with representative support-expanded evaluations: (left) we train on a black background and evaluate on cumulative background-color supports (black, black+white, black+white+red, black+white+red+green, black+white+red+green+blue); (right) we train with no distractors and evaluate with increasing numbers of same-as-agent distractors (0, 1, 2, 3, 5, 7, 9, 11), where distractors match the agent’s shape and color. Across suites, training success rises rapidly, while evaluation success often saturates substantially lower, revealing persistent train–eval gaps under purely visual shifts with fixed dynamics and rewards.

Generalization is strongly axis-dependent (mean\(\pm\)SEM). Ranking axes by success-rate degradation, the largest gaps arise from filters (\(\Delta\)SR \(=86.8\%\)) and effects (\(80.5\%\)), followed by layout (\(62.8\%\)) and background (\(53.3\%\)); distractors (\(30.9\%\)) and agent appearance (\(21.1\%\)) are comparatively milder (1).

Background shifts impair both motion and completion. Averaged across background pairs, distance and progress drop by \(30.5\%\) and SR drops from \(0.90\) to \(0.42\) (\(\Delta\)SR \(=53.3\%\)), accompanied by a large absolute return gap. In 6 (left), evaluation success decreases as the evaluation support expands from the training background to mixtures that include unseen background colors, while training success on the black background remains high. These curves should be read as support-expansion tests: for example, the black+white+red+green setting evaluates a mixture over those colors, so performance reflects success on the training background and failures on newly included backgrounds. Concretely, support expansion is implemented across evaluation episodes, not by blending backgrounds within an episode: black+white assigns half of evaluation episodes to each color, while black+white+red assigns one third of episodes to each color.

Photometric and lighting perturbations primarily break completion. For filters and effects, distance degradation is moderate (\(\Delta\)Dist \(=11.7\%\) and \(20.8\%\)), yet SR collapses (\(0.83\!\rightarrow\!0.11\) and \(0.82\!\rightarrow\!0.16\); \(\Delta\)SR \(=86.8\%\) and \(80.5\%\)), indicating that motion and shaped reward can persist while success fails under photometric/lighting shifts.

Small motion gaps can mask large completion gaps. Distractors and layout show small distance/progress gaps (\(\sim\)​3–4%) but sizable SR drops (\(30.9\%\) and \(62.8\%\)). In 6 (right), increasing same-as-agent distractors (0–11) progressively suppresses evaluation success with unchanged training success.

Per-configuration behavior is heterogeneous. 2 includes both negligible-gap sanity checks and near-failure pairs, e.g., black\(\rightarrow\)noise backgrounds (\(\Delta\)SR \(=98.9\%\)), hue shift \(180^\circ\) (\(98.8\%\)), light count \(4\) (\(95.5\%\)), and 7 same-as-agent distractors (\(92.0\%\)). Within Background, training with more visual diversity reduces SR gaps. 13 provides full learning curves for all 34 pairs; some small return gaps arise because both train and eval fail, motivating joint reporting of distance, progress, and SR. Overall, PPO-CNN is strong in-distribution but brittle under controlled visual shifts, with failures concentrated in task completion rather than basic locomotion. The additional robustness baselines in 12, especially 4, support the same conclusion. RAD-PPO with cutout-color improves over PPO-CNN on 7 of the 8 hardest configurations, but substantial success-rate gaps remain. In contrast, RAD-PPO with crop-resize augmentation and off-policy CURL, DrQ, and SVEA did not converge in these platformer runs. Because the non-geometric RAD-PPO cutout-color variant converged as expected, we attribute these failures to spatial augmentations disrupting screen positions needed for cliff distance estimation, obstacle avoidance, and jump timing.

8 Limitations↩︎

KAGE-Bench is intentionally scoped as a diagnostic benchmark for known-axis visual distribution shifts, not as a comprehensive proxy for real-world RL. The current suite uses a single 2D platformer with fixed dynamics and rewards so that observed performance changes can be cleanly attributed to visual factors. The visual shifts are stylized and factorized rather than photorealistic, and the reported suites were selected using a representative PPO-CNN probe; stronger or differently structured agents may induce different axis rankings. Consequently, the ordering of axes should be interpreted as a property of this controlled benchmark and baseline family, not as a universal ranking of visual robustness difficulty. Future extensions should include geometric and perspective shifts, richer task families, and complementary benchmarks where controlled attribution is traded for broader realism. The environment is simple at the latent-dynamics level by design, but this does not remove all spatial demands: near-complete traversal still requires consistent pixel-to-layout correspondence and accurate jump timing. Adding mechanics such as double jumps, obstacles, or richer objectives is supported by the configurable environment interface, but such additions should be treated as new benchmark variants because they introduce interacting factors beyond isolated visual shift. Continual and curriculum learning are also natural uses of KAGE-Env, but they are outside the scope of the present benchmark.

9 Conclusion↩︎

We introduced KAGE-Env, a JAX-native RL environment for controlled studies of visual generalization that factorizes the observation process into independently configurable visual axes while keeping the underlying control problem fixed, enabling high-throughput evaluation via end-to-end compilation and large-scale parallel simulation. Building on this environment, we presented KAGE-Bench, a standardized benchmark comprising six known-axis suites and 34 train–evaluation configuration pairs that isolate specific sources of visual shift and allow precise attribution of performance changes. Empirically, we find that visual generalization difficulty varies substantially across axes: background changes and photometric or lighting perturbations induce the most severe failures, often collapsing task success despite nontrivial progress, whereas agent-appearance shifts are comparatively benign. Overall, KAGE-Bench provides a fast, reproducible, and diagnostic framework for evaluating pixel-based RL under controlled visual variation, and we expect it to support more systematic analysis of visual robustness and future work on richer shifts, broader task families, and alternative learning algorithms.

Impact Statement↩︎

This paper presents work whose goal is to advance the field of machine learning by introducing a fast, reproducible benchmark for studying visual generalization in RL. We do not anticipate immediate negative societal impacts from releasing an evaluation environment and configuration suites; however, as with most progress in robust perception and control, improved generalization methods could enable more capable autonomous systems, which may have downstream applications with safety and misuse considerations. We hope KAGE-Env and KAGE-Bench support more rigorous and transparent evaluation of robustness, helping the community identify failure modes early and develop safer learning systems.

10 Reducing Visual Shifts to State-Policy Shifts↩︎

10.1 Problem setup.↩︎

KAGE-Bench is constructed to isolate purely visual distribution shift. Formally, each environment instance is indexed by a visual configuration \(\xi\in\Xi\) (e.g., the YAML parameters controlling background, filters, lighting, sprites), and \(\xi\) determines how a latent simulator state \(s\in S\) is rendered into a pixel observation \(o\in\Omega\). This rendering mechanism is modeled as an observation kernel \(O_\xi(\cdot\mid s)\), meaning that, given the same latent state \(s\), different \(\xi\) may produce different distributions over images. Crucially, KAGE-Bench enforces that \(\xi\) does not alter the control problem itself: the transition kernel \(P(\cdot\mid s,a)\) and reward function \(r(s,a)\) are identical for all \(\xi\). Hence, when we observe a train–test gap after changing \(\xi\), it cannot be caused by different dynamics or rewards; it must be caused by the interaction between the same observation-based policy and a different rendering process.

The key point is that a policy trained on pixels, \(\pi(a\mid o)\), does not directly specify actions as a function of the latent state \(s\), but only as a function of the rendered image \(o\). Therefore, the action distribution conditioned on the latent state depends on \(\xi\) through the distribution of renderings \(O_\xi(\cdot\mid s)\). The definition below formalizes this dependence by defining, for each \(\xi\), an induced state policy (2):

Induced state policy is the conditional distribution of the action after integrating out (marginalizing) the intermediate observation variable \(o\). In this sense, changing \(\xi\) is equivalent to changing the induced state policy: even if the pixel policy \(\pi(a\mid o)\) is fixed, the effective mapping from latent states to action distributions changes because the policy is evaluated on different renderings. This reduction is fundamental for analysis: it converts visual generalization under observation shifts into a standard policy shift problem in a fixed latent MDP.

10.2 Setting and central objects↩︎

We work with the following measurable objects.

• Measurable spaces. \((S,\mathcal{S})\) is the latent state space equipped with a \(\sigma\)-algebra \(\mathcal{S}\) of measurable subsets; \((A,\mathcal{A})\) is the action space equipped with \(\sigma\)-algebra \(\mathcal{A}\); \((\Omega,\mathcal{O})\) is the observation (pixel). Measurability ensures that probabilities and integrals used below are well-defined.

• MDP primitives. \(\gamma\in[0,1)\) is the discount factor and \(\rho_0\) is the initial distribution on \((S,\mathcal{S})\).

• Transition kernel (Markov kernel). \(P(\cdot \mid s,a)\) specifies the environment dynamics. For every state–action pair \((s,a)\in S\times A\), \(P(\cdot \mid s,a)\) is a probability distribution over next states in \(S\). Operationally, this means that after taking action \(a_t\) in state \(s_t\), the next state is sampled as \(s_{t+1} \sim P(\cdot \mid s_t,a_t)\).

• Reward function. \(r:S\times A\to\mathbb{R}\) is measurable and bounded: \(\|r\|_\infty := \sup_{(s,a)\in S\times A} |r(s,a)| < \infty\). Boundedness guarantees the discounted return \(\sum_{t\ge 0}\gamma^t r(s_t,a_t)\) is integrable.

• Visual configuration space. \(\Xi\) indexes renderers. For each \(\xi\in\Xi\), \(O_\xi(\cdot\mid s)\) is an observation kernel (a Markov kernel from \((S,\mathcal{S})\) to \((\Omega,\mathcal{O})\)). Operationally, given latent state \(s\), an image is sampled as \(o\sim O_\xi(\cdot\mid s)\).

• Reactive pixel policy. \(\pi(\cdot\mid o)\) is a Markov kernel from \((\Omega,\mathcal{O})\) to \((A,\mathcal{A})\) (memoryless policy): given observation \(o\), an action is sampled as \(a\sim \pi(\cdot\mid o)\).

10.2.0.1 Latent MDP and visual POMDP.

10.3 Main theorem.↩︎

10.4 Proof of Theorem 2↩︎

10.4.0.1 Step 0 (Generative dynamics in \(\mathcal{M}_\xi\)).

By definition of the POMDP \(\mathcal{M}_\xi\) (4) and the reactive policy \(\pi\), the interaction at time \(t\) is:

\[\label{eq:gen95obs} \left\{ \begin{align} &o_t \sim O_\xi(\cdot\mid s_t) \\ &a_t \sim \pi(\cdot\mid o_t) \\ &s_{t+1} \sim P(\cdot\mid s_t,a_t) \end{align} \right.\tag{4}\]

Equation 4 (top) represents the rendering step: it formalizes that pixels are generated from the latent state via \(O_\xi\). Equation 4 (middle) is the policy step: the agent samples an action using only the pixels. Equation 4 (bottom) is the environment dynamics: the next state depends only on \((s_t,a_t)\) through \(P\) and is independent of \(o_t\) given \((s_t,a_t)\). Therefore, observations influence the future only through their effect on the chosen action.

10.4.0.2 Step 1 (Show conditional action law 2 ).

Fix a time \(t\ge 0\) and an arbitrary measurable set \(B\in\mathcal{A}\). We compute \(\mathbb{P}(a_t\in B\mid s_t)\) by conditioning on the intermediate variable \(o_t\) (the observation).

(1a) Law of total probability (tower property). Recall the tower property of conditional expectation [49]: for any integrable random variable \(X\) and \(\sigma\)-algebras \(\mathcal{G}_1 \subseteq \mathcal{G}_2\), \[\label{eq:tower95general} \mathbb{E}[X \mid \mathcal{G}_1] = \mathbb{E}\!\left[ \mathbb{E}[X \mid \mathcal{G}_2] \,\middle|\, \mathcal{G}_1 \right].\tag{5}\] This identity states that conditioning can be performed in stages: one may first condition on a finer information set \(\mathcal{G}_2\) and then average again while conditioning on the coarser information set \(\mathcal{G}_1\).

We apply 5 to the indicator random variable \[X := \mathbb{I}\{a_t \in B\},\] which is integrable since it is bounded between \(0\) and \(1\). Recall that conditional probabilities can be written as conditional expectations of indicator functions: \[\mathbb{P}(a_t \in B \mid \mathcal{G}) = \mathbb{E}\!\left[\mathbb{I}\{a_t \in B\} \mid \mathcal{G}\right].\]

Next, we specify the two \(\sigma\)-algebras:

• \(\mathcal{G}_1 := \sigma(s_t)\), the \(\sigma\)-algebra generated by the latent state \(s_t\) (i.e., conditioning on knowing \(s_t\)),

• \(\mathcal{G}_2 := \sigma(o_t, s_t)\), the \(\sigma\)-algebra generated by the pair \((o_t, s_t)\) (i.e., conditioning on knowing both the observation and the latent state).

Clearly, \(\mathcal{G}_1 \subseteq \mathcal{G}_2\), since knowing \((o_t,s_t)\) includes knowing \(s_t\).

Applying 5 with these choices gives \[\begin{align} \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t \in B \mid s_t) &= \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \mathbb{I}\{a_t \in B\} \,\middle|\, s_t \right] \\ &= \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \mathbb{I}\{a_t \in B\} \,\middle|\, o_t, s_t \right] \,\middle|\, s_t \right]. \end{align}\]

Finally, rewriting the inner conditional expectation again as a conditional probability yields \[\label{eq:tower} \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t\in B\mid s_t) = \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t\in B\mid o_t,s_t) \,\middle|\, s_t \right].\tag{6}\]

This equality formalizes the intuitive idea that, to compute the probability of choosing an action in \(B\) given the latent state \(s_t\), one may first compute this probability given the more detailed information \((o_t,s_t)\) and then average over all possible observations \(o_t\) that can occur when the state is \(s_t\).

(1b) Use the policy sampling rule. Recall from the interaction dynamics that, at time \(t\), once the observation \(o_t\) is generated, the action is sampled according to the policy: \[a_t \sim \pi(\cdot \mid o_t).\] This means that the conditional distribution of \(a_t\) given \(o_t\) is exactly \(\pi(\cdot \mid o_t)\).

Formally, for any measurable action set \(B \in \mathcal{A}\), \[\mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t \in B \mid o_t) = \pi(B \mid o_t).\]

Moreover, because the policy is reactive (memoryless), the action depends on the current observation \(o_t\) but not directly on the latent state \(s_t\) once \(o_t\) is known. Therefore, conditioning additionally on \(s_t\) does not change the conditional distribution: \[\label{eq:policy95conditional} \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t \in B \mid o_t, s_t) = \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t \in B \mid o_t) = \pi(B \mid o_t).\tag{7}\]

This equality expresses the fact that the policy fully mediates the influence of the observation on the action, and no additional information about \(s_t\) is used once \(o_t\) has been observed.

(1c) Substitute 7 into 6 . \[\label{eq:substitute1} \mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t\in B\mid s_t) = \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \pi(B\mid o_t) \,\middle|\, s_t \right].\tag{8}\] At this point, the only remaining randomness inside the conditional expectation comes from \(o_t\) given \(s_t\).

(1d) Use the observation sampling rule. Since \(o_t\mid s_t \sim O_\xi(\cdot\mid s_t)\) by 4 , the conditional expectation in 8 can be written as an integral with respect to the measure \(O_\xi(\cdot\mid s_t)\): \[\label{eq:substitute2} \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \pi(B\mid o_t) \,\middle|\, s_t \right] = \int_{\Omega} \pi(B\mid o)\, O_\xi(do\mid s_t).\tag{9}\] This step is precisely what “averaging over renderings” means: we are averaging the policy’s probability of selecting an action in \(B\) over all images \(o\) that can be rendered from \(s_t\) under configuration \(\xi\).

(1d) Use the observation sampling rule. At this point, the random quantity inside the conditional expectation in 8 is \(\pi(B\mid o_t)\), and the only remaining source of randomness is the observation \(o_t\) given the latent state \(s_t\). By the generative dynamics of the visual POMDP 4 , the observation at time \(t\) is sampled according to the observation kernel: \(o_t \mid s_t \sim O_\xi(\cdot \mid s_t)\). Therefore, conditioning on \(s_t\), the random variable \(\pi(B\mid o_t)\) is distributed according to the pushforward of \(O_\xi(\cdot\mid s_t)\) through the function \(o \mapsto \pi(B\mid o)\).

By the definition of conditional expectation with respect to a Markov kernel, this conditional expectation can be written as an integral over the observation space: \[\label{eq:substitute3} \mathbb{E}_{\mathcal{M}_\xi,\pi}\!\left[ \pi(B\mid o_t) \,\middle|\, s_t \right] = \int_{\Omega} \pi(B\mid o)\, O_\xi(do\mid s_t).\tag{10}\]

This expression makes explicit what is meant by “averaging over renderings”: for a fixed latent state \(s_t\), we take all images \(o\) that the renderer may produce under configuration \(\xi\), weight the policy’s action probability \(\pi(B\mid o)\) by how likely each image is under \(O_\xi(\cdot\mid s_t)\), and sum (integrate) these contributions. The result is the average probability of selecting an action in \(B\) after accounting for all possible renderings of the same latent state.

(1e) Recognize the induced policy definition. By 2, the right-hand side of 10 equals \(\pi_\xi(B\mid s_t)\). Therefore, \[\mathbb{P}_{\mathcal{M}_\xi,\pi}(a_t\in B\mid s_t) = \pi_\xi(B\mid s_t),\] which is exactly 2 . This completes Item 1.

10.4.0.3 Step 2 (Equality in law of state–action processes.).

We now show that the state–action process in \(\mathcal{M}_\xi\) under \(\pi\) evolves exactly as in the latent MDP \(\mathcal{M}\) under \(\pi_\xi\).

(2a) Effective action selection given \(s_t\). Item 1 implies that, conditional on \(s_t\), the action \(a_t\) has distribution \(\pi_\xi(\cdot\mid s_t)\). Hence, if we are interested only in the joint process \((s_t,a_t)\) (and not in \(o_t\)), we may replace the two-step procedure \[o_t \sim O_\xi(\cdot\mid s_t),\qquad a_t\sim \pi(\cdot\mid o_t)\] by the single step \[a_t \sim \pi_\xi(\cdot\mid s_t),\] without changing the conditional distribution of \(a_t\) given \(s_t\).

(2b) State transition given \((s_t,a_t)\) is identical. Under \(\mathcal{M}_\xi\), the next state satisfies \(s_{t+1}\sim P(\cdot\mid s_t,a_t)\) by 4 . This is exactly the same transition rule as in the latent MDP \(\mathcal{M}\), and it depends only on \((s_t,a_t)\).

(2c) Conclude identical recursion. Combining (2a) and (2b), the pair \((s_t,a_t)\) evolves according to \[s_0\sim\rho_0,\qquad a_t\sim \pi_\xi(\cdot\mid s_t),\qquad s_{t+1}\sim P(\cdot\mid s_t,a_t).\] This is precisely the generative definition of executing the state policy \(\pi_\xi\) in the latent MDP \(\mathcal{M}\). Therefore, the joint laws of \((s_0,a_0,s_1,a_1,\ldots)\) coincide under \((\mathcal{M}_\xi,\pi)\) and \((\mathcal{M},\pi_\xi)\), proving Item 2.

10.4.0.4 Step 3 (Equality of expected discounted return).

Define the discounted return random variable \[G := \sum_{t=0}^\infty \gamma^t r(s_t,a_t).\]

We first verify that \(G\) is integrable. By assumption, the reward function \(r\) is bounded, meaning that for all \((s,a)\in S\times A\), \[|r(s,a)| \le \|r\|_\infty < \infty.\] Therefore, for every time step \(t\), \[|\gamma^t r(s_t,a_t)| \le \gamma^t \|r\|_\infty.\]

Summing these bounds over \(t\) and using that \(\gamma\in[0,1)\) yields \[|G| = \left|\sum_{t=0}^\infty \gamma^t r(s_t,a_t)\right| \le \sum_{t=0}^\infty \gamma^t |r(s_t,a_t)| \le \sum_{t=0}^\infty \gamma^t \|r\|_\infty.\]

The right-hand side is a convergent geometric series: \[\sum_{t=0}^\infty \gamma^t \|r\|_\infty = \|r\|_\infty \sum_{t=0}^\infty \gamma^t = \frac{\|r\|_\infty}{1-\gamma} < \infty.\] Hence, \(G\) is almost surely finite and integrable.

By Item 2, the state–action processes \((s_t,a_t)_{t\ge 0}\) have the same law under \((\mathcal{M}_\xi,\pi)\) and \((\mathcal{M},\pi_\xi)\). Since \(G\) is a measurable function of the entire state–action trajectory and depends only on \((s_t,a_t)\), it follows that \(G\) has the same distribution under both constructions. In particular, their expectations coincide: \[J(\pi;\mathcal{M}_\xi) = \mathbb{E}_{\mathcal{M}_\xi,\pi}[G] = \mathbb{E}_{\mathcal{M},\pi_\xi}[G] = J(\pi_\xi;\mathcal{M}).\] This proves Item 3 and completes the proof. 0◻

10.5 Interpretation for KAGE-Bench↩︎

2 provides a precise formal justification for how visual generalization should be interpreted in KAGE-Bench. Because the latent dynamics \(P\) and reward function \(r\) are identical across all visual configurations \(\xi\), the theorem shows that changing \(\xi\) affects the learning problem only through the observation channel \(O_\xi(\cdot\mid s)\). For any fixed pixel policy \(\pi(a\mid o)\), this change manifests exclusively as a change in the induced state-conditional action distribution \(\pi_\xi(\cdot\mid s)\).

Crucially, the theorem establishes an exact equivalence in distribution at the level of latent state–action trajectories: executing the observation-based policy \(\pi\) in the visual POMDP \(\mathcal{M}_\xi\) produces the same joint law over \((s_t,a_t)\) as executing the induced state policy \(\pi_\xi\) in the latent MDP \(\mathcal{M}\). This result is purely representational. It does not claim that \(\pi_\xi\) is optimal, nor that marginalizing over observations improves performance. Rather, it shows that all effects of visual variation are captured entirely by the induced policy, without altering the underlying control problem.

This equivalence is central to the design and interpretation of KAGE-Bench. It guarantees that any observed train–test performance gap under a visual shift \(\xi \to \xi'\) cannot be attributed to changes in dynamics, rewards, or task structure, but must correspond exactly to a performance difference between two state policies \(\pi_\xi\) and \(\pi_{\xi'}\) acting in the same latent MDP. As a consequence, KAGE-Bench reduces visual generalization to a well-defined policy shift problem in a fixed MDP, enabling principled analysis using standard reinforcement learning tools and ensuring that benchmark results isolate perception-induced failures rather than confounding control effects.

10.6 Additional consequences: equivalence of trajectory-level metrics↩︎

2 implies more than equality of expected return. Because it establishes equality in distribution of the latent state–action process \((s_t,a_t)_{t\ge 0}\), any performance metric that is a measurable function of the latent trajectory inherits the same equivalence. We formalize this as a corollary.

10.6.0.1 Proof.

By Item 2 of Theorem 2, the joint laws of the state–action trajectories coincide: \[\mathcal{L}_{\mathcal{M}_\xi,\pi}\big((s_t,a_t)_{t\ge 0}\big) = \mathcal{L}_{\mathcal{M},\pi_\xi}\big((\tilde{s}_t,\tilde{a}_t)_{t\ge 0}\big).\] Applying any measurable function \(F\) to two random elements with the same law yields random variables with the same law. Equality of expectations follows immediately. 0◻

We now specialize Corollary 3 to the concrete evaluation metrics used in KAGE-Bench.

10.6.0.2 Interpretation.

Corollary 4 shows that the equivalence established in Theorem 2 applies not only to discounted return, but also to all trajectory-based evaluation metrics commonly reported in KAGE-Bench, including raw distance traveled, normalized progress, and binary success. These quantities depend only on the latent state trajectory \((s_t)_{t\ge 0}\) and are therefore fully determined by the induced state policy \(\pi_\xi\) in the latent MDP.

As a result, differences in success rate, progress, or distance under a visual shift \(\xi\to\xi'\) are exactly differences between the induced state policies \(\pi_\xi\) and \(\pi_{\xi'}\) acting in the same latent MDP. This further reinforces that KAGE-Bench isolates perception-induced failures: all reported metrics admit a clean interpretation as properties of state-policy shift rather than changes in the underlying control task.

11 Extended figures and tables↩︎

7 compares training success rate with evaluation success rate under increasingly broad variants of a single visual axis. 3 shows train and eval results across each reported metric across each config.

12 Additional robustness baseline results↩︎

This appendix preserves the additional experiments introduced during rebuttal. The table compares PPO-CNN with RAD-PPO using cutout-color augmentation on the hardest KAGE-Bench configurations, and the figures visualize why spatial augmentations such as crop-resize and random shift can damage platformer observations.

3.5pt

13 Benchmark Training Details↩︎

This appendix reports the full learning curves for the PPO-CNN baseline on all \(34\) KAGE-Bench train-evaluation configuration pairs. At each logging checkpoint, we evaluate the current policy on both the corresponding training configuration (in-distribution) and its paired evaluation configuration (out-of-distribution), and plot the resulting metrics over environment steps. Figures are grouped by generalization axis: Agent Appearance (11); Background (12, 13); Distractors (15); Effects (16); Filters (17, 18); and Layout (14).

14 Generalization axes review↩︎

15 Backgrounds↩︎

16 Agents Sprites↩︎

17 YAML Configuration Details↩︎

# ==============================================================================
# JAX Platformer Environment Configuration
# ==============================================================================
# This configuration file allows precise control over all environment aspects.
# Edit this file to customize backgrounds, characters, NPCs, physics, and usage.
#
# Load this config using:
#   from kage_bench import EnvConfig, load_config_from_yaml
#   config = load_config_from_yaml("custom_config.yaml")
#   env = PlatformerEnv(config)
# ==============================================================================

# ------------------------------------------------------------------------------
# 0. Episode Settings
# ------------------------------------------------------------------------------
episode_length: 500        # Default: 500 | Max steps per episode
forward_reward_scale: 0.2  # Default: 0.2 | Reward for moving forward
jump_penalty: 10.0         # Default: 10.0 | Penalty for jumping
timestep_penalty: 0.1      # Default: 0.1 | Per-timestep reward penalty
idle_penalty: 5.0          # Default: 5.0 | Penalty when x does not change
dist_to_success: 490.0     # Default: 490.0 | Passed distance needed for success


# ------------------------------------------------------------------------------
# 1. Global Screen Settings
# ------------------------------------------------------------------------------
H: 128  # Default: 128 | Screen height in pixels
W: 128  # Default: 128 | Screen width in pixels

# 2. Background Settings
# ------------------------------------------------------------------------------
background:
  # * Mode (Default: "black"):
    # "black" - just black background
    # "image" - image backgrounds
    # "noise" - white noise background (unique for each episode)
    # "color" - color backgrounds
  mode: "image"

  # * Image mode settings (ignired if mode != "image"):
  # * Please, choose only one of the following options image_dir, image_paths, image_path:
  # [Option 1] Specify a directory to load all images from (randomly selected per episode)
  image_dir: "src/kage_bench/assets/backgrounds" # * Will use all the 128 images

  # [Option 2] Or specify list of explicit paths:
  # image_paths: # * Will use only listed images
  #   - "src/kage_bench/assets/backgrounds/bg-1.jpeg"
  #   - "src/kage_bench/assets/backgrounds/bg-64.jpeg"
  #   - "src/kage_bench/assets/backgrounds/bg-128.jpeg"

  # [Option 3] Or force a single image:
  # image_path: "path/to/your/image.jpeg"

  # Parallax/Tiling (ignored if mode != "image")
  parallax_factor: 0.5   # Default: 0.5 | 0.0 = static, 0.5 = slow scroll, 1.0 = locked to camera, <0 = moves left, >1 = moves fast
  tile_horizontal: true  # Default: true | Repeat image horizontally for infinite worlds

  # Dynamic switching (change background during the episode \;gif effect)
  switch_frequency: 0.0  # Default: 0.0 | Probability per step (0.0 = never switch, 1.0 = every step)

  # * Color mode settings (ignored if mode != "color"):
  # List of colors to randomly select from per episode
  # Available colors:
    # "black", "white", "red", "orange", "yellow", "green", "cyan", "blue",
    # "purple", "pink", "brown", "gray", "lime", "teal", "indigo", "magenta"
  color_names: ["purple", "teal", "indigo"]

# ------------------------------------------------------------------------------
# 3. Character Settings (Player)
# ------------------------------------------------------------------------------
character:
  width: 16  # Default: 16 | Character width in pixels  | Recommended to not change
  height: 24 # Default: 24 | Character height in pixels | Recommented to not change

  # * Sprite Settings (ignored if use_shape: true):
  use_sprites: true # Default: true | Use sprites for character
  # * Please, choose only one of the following options sprite_dir, sprite_paths, sprite_path:
  # [Option 1] Directory containing sprite subdirectories (auto-discovers all subdirs as skins):
  sprite_dir: "src/kage_bench/assets/sprites"  # Will auto-discover and use all subdirectories like clown/, robot/, skelet/, etc.

  # [Option 2] List of directories for multiple skins (randomly selected per episode):
  # NOTE: Each directory should contain .png animation frames (e.g., walk-1.png, walk-2.png)
  # sprite_paths:
    # - "src/kage_bench/assets/sprites/clown"         # Skin 1
    # - "src/kage_bench/assets/sprites/dark_knight"   # Skin 2
    # - "src/kage_bench/assets/sprites/skeleton"        # Skin 3

  # [Option 3] Single sprite directory:
  # sprite_path: "path/to/your/folder/with/sprite"

  # Animation settings (ignored if use_sprites: false):
  enable_animation: true  # Default: true | Enable sprite animation (if false, use only first frame)
  animation_fps: 12.0     # Default: 12.0 | Frames per second for sprite animation
  idle_sprite_idx: 0      # Default: 0 | Which sprite to use when idle (typically first)


  # * Shape Settings (ignored if use_sprites: true):
  use_shape: false # Default: false | Use shapes for character

  # Available shapes:
    # "circle", "cross", "diamond", "ellipse", "line",
    # "polygon", "square", "star", "triangle"
  shape_types: ["circle", "star"]

  # Available colors:
    # "red", "green", "blue", "orange", "yellow", "violet", "magenta",
    #"cyan", "pink", "brown", "purple", "lime", "navy", "maroon",
    # "olive", "teal", "indigo", "coral", "gold", "silver", "white"
  shape_colors: ["teal", "indigo"]
  shape_rotate: true        # Default: true | Enable shape rotation
  shape_rotation_speed: 5.0 # Default: 5.0 | Rotation speed in degrees per second

# ------------------------------------------------------------------------------
# 4. NPC Settings (Non-Player Characters)
# ------------------------------------------------------------------------------
npc:
  # Default: true | Enable NPC system:
  # * World-Fixed NPCs (stand on platforms)
  enabled: true       # Default: true | Enable World-Fixed NPC system:
  min_npc_count: 5    # Default: 5 | Minimum number of NPCs per level
  max_npc_count: 20   # Default: 20 | Maximum number of NPCs per level
  spawn_y_offset: 0   # Default: 0 | Offset from ground in pixels (positive = up, negative = down)
  animation_fps: 12.0 # Default: 12.0 | Animation speed for NPCs

  # * Please, choose only one of the following options sprite_dir, sprite_paths, sprite_path:
  # [Option 1] Specify a directory to auto-discover all sprite subdirectories:
  sprite_dir: "src/kage_bench/assets/sprites"  # Will auto-discover all subdirs like robot/, girl/, skelet/, etc.

  # [Option 2] Or specify list of explicit paths:
  # sprite_paths:
  #   - "src/kage_bench/assets/sprites/clown"         # Skin 1
  #   - "src/kage_bench/assets/sprites/dark_knight"   # Skin 2
  #   - "src/kage_bench/assets/sprites/skeleton"        # Skin 3

  # [Option 3] Or force a single sprite directory:
  # sprite_path: "path/to/your/folder/with/sprite"

  # * Sticky NPCs (follow camera and always visible in observation):
  sticky_enabled: false  # Default: false | Enable sticky NPC system:
  min_sticky_count: 1    # Default: 1 | Minimum number of sticky NPCs per level
  max_sticky_count: 5    # Default: 5 | Maximum number of sticky NPCs per level
  # * Please, choose only one of the following options sticky_sprite_dir, sticky_sprite_dirs, sticky_sprite_path:
  # [Option 1] Specify a directory to auto-discover all sticky sprite subdirectories:
  sticky_sprite_dir: "src/kage_bench/assets/sprites"  # Will auto-discover all subdirs

  # [Option 2] Or specify list of explicit paths (backward compatibility):
  # sticky_sprite_dirs: # * Will use only listed sprite directories
  #   - "src/kage_bench/assets/sprites/clown"         # Skin 1
  #   - "src/kage_bench/assets/sprites/dark_knight"   # Skin 2
  #   - "src/kage_bench/assets/sprites/skeleton"        # Skin 3

  # [Option 3] Or force a single sticky sprite directory:
  # sticky_sprite_path: "path/to/your/folder/with/sprite"

  # Sticky NPCs settings:
  sticky_can_jump: true          # Default: true | Enable jumping for sticky NPCs
  sticky_jump_probability: 0.01  # Default: 0.01 | Probability of jumping per step
  sticky_y_min_offset: -40       # Default: -40  | Minimum Y offset from ground (negative = higher)
  sticky_y_max_offset: -10       # Default: -10  | Maximum Y offset from ground (0 = on ground)
  sticky_x_offsets: []           # Default: []   | Camera-relative X offsets for sticky NPCs (e.g., [-40, 0, 40]). If not provided, NPCs will be spread around agent
  sticky_x_min: -60              # Default: -60  | Minimum X offset from center
  sticky_x_max: 60               # Default: 60   | Maximum X offset from center

# ------------------------------------------------------------------------------
# 5. Distractors Settings
# ------------------------------------------------------------------------------
distractors:
  # Whether to enable distractors, i.e. moving geometric shapes on background (always visible in observation)
  enabled: true # Default: false | Enable distractors system:
  count: 5      # Default: 5 | Number of distractors per level

  # Available shapes:
    # "circle", "cross", "diamond", "ellipse",
    # "line", "polygon", "square", "star", "triangle"
  shape_types: ["circle", "star", "cross"]

  # Available colors:
    # "red", "green", "blue", "orange", "yellow", "violet", "magenta",
    # "cyan", "pink", "brown",  "purple", "lime", "navy", "maroon",
    # "olive", "teal", "indigo", "coral", "gold", "silver", "white"
  shape_colors: ["red", "green", "blue"]

  # Dynamics:
  can_move: true           # Default: true | Whether distractors move around
  min_speed: 0.0           # Default: 0.0 | Minimum movement speed (pixels per step)
  max_speed: 1.0           # Default: 2.0 | Maximum movement speed (pixels per step)
  can_rotate: true         # Default: true | Whether distractors rotate
  min_rotation_speed: -0.3 # Default: -3.0 | Minimum rotation speed (degrees per step)
  max_rotation_speed: 0.3  # Default: 3.0 | Maximum rotation speed (degrees per step)
  min_size: 4              # Default: 4 | Minimum size
  max_size: 12             # Default: 12 | Maximum size

# ------------------------------------------------------------------------------
# 6. Filter & Effect Settings
# ------------------------------------------------------------------------------
filters:
  brightness: 0.0        # Default: 0.0 | [-1, 1] additive exposure, scaled in code
  contrast: 1.0          # Default: 1.0 | >0, scales around mid-gray (128)
  gamma: 1.0             # Default: 1.0 | [0.5, 2.0] power-law on [0,1]
  saturation: 1.0        # Default: 1.0 | [0, 2] HSV S multiplier
  hue_shift: 0.0         # Default: 0.0 | [-180, 180] degrees, HSV hue offset
  color_temp: 0.0        # Default: 0.0 | [-1, 1] warm(+R,-B) vs cool(+B,-R)

  # Stochastic effects (require PRNG key):
  color_jitter_std: 0.0    # Default: 0.0 | >=0, std of 3x3 RGB mixing perturbation
  gaussian_noise_std: 0.0  # Default: 0.0 | >=0, pixelwise N(0, std^2) in [0,255]
  poisson_noise_scale: 0.0 # Default: 0.0 | [0,1], shot noise with lambda = img*scale

  # Spatial / detail transforms:
  blur_sigma: 0.0        # Default: 0.0 | >=0, box-blur approximation strength
  sharpen_amount: 0.0    # Default: 0.0 | >=0, unsharp mask gain
  pixelate_factor: 1     # Default: 1 | int>=1, down/up nearest (1 disables)

  # Global shading / lighting overlays:
  vignette_strength: 0.0     # Default: 0.0 | >=0, edge darkening factor (code expects ~[0,1])
  radial_light_strength: 0.0 # Default: 0.0 | >=0, additive center light (code expects ~[0,1])

  # Optional preset stack:
  pop_filter_list: []    # Default: [] | ["vintage","retro","cyberpunk","horror","noir"]

# ------------------------------------------------------------------------------
# 7. Effects Settings
# ------------------------------------------------------------------------------
effects:
  # Point light effect
  point_light_enabled: false  # Default: false | Enable/disable point light effects
  point_light_intensity: 1.0  # Default: 1.0 | Light intensity in [0.1, 5.0]
  point_light_radius: 0.1     # Default: 0.1 | Light radius as fraction of image size [0.01, 1.0]
  point_light_falloff: 2      # Default: 2.0 | Falloff exponent in [1.0, 4.0], higher = sharper

  # Multiple random lights
  point_light_count: 4        # Default: 1 | Number of lights in [1, 5]

  # Available colors:
    # `warm_white`, `cool_white`, `yellow`, `orange`, `red`,
    # `green`, `cyan`, `blue`, `purple`, `pink`, `gold`, `fire`
  point_light_color_names: ["blue", "pink", "gold"]   # Default: ["warm_white"] | List of color names from LIGHT_COLORS

# ------------------------------------------------------------------------------
# 8. Level Layout
# ------------------------------------------------------------------------------
layout:
  length: 2048            # Default: 2048 | Length of the level in pixels
  height_px: 128          # Default: 128 | Height of the level in pixels
  base_ground_y: 96       # Default: 96 | [70, 127] Hight of the ground level. Higher - lower ground level
  pix_per_unit: 2         # Default: 2 | [0, 3] Pixels per height unit. 0 - flat, 3 - larger steps
  ground_thickness: 2     # Default: 2 | [1, 10] Thickness of the ground band in pixels
  run_width: 25           # Default: 20 | [1, 60] Widths of the stair in pixels
  p_change: 0.7           # Default: 0.7 | [0, 1] Probability of height change. 0 -never change, 1 -always change
  p_up_given_change: 0.5  # Default: 0.5 | [0, 1] Probability of height increase given height change. 0 -always decrease, 1 -always increase
  min_step_height: 5      # Default: 5 | [1, 17] Minimum height of the step in pixels
  max_step_height: 17     # Default: 10 | [1, 17] Maximum height of the step in pixels

  # Colors for platforms (randomly selected per episode). Available colors:
    # "black", "white", "red", "orange", "yellow", "green", "cyan", "blue",
    # "purple", "pink", "brown", "gray", "lime", "teal", "indigo", "magenta"
  layout_colors: ["cyan"] # Default: ["cyan"] | List of color names from COLOR_PALETTE


# ------------------------------------------------------------------------------
# 9. Physics Settings
# ------------------------------------------------------------------------------
physics:
  gravity: 0.75         # Default: 0.75 | [0.1, 1.0] Gravity force
  move_speed: 1         # Default: 1 | int and > 1 | Move speed
  jump_force: -7.5      # Default: -7.5 | [-10.0, 0.0] Jump force. Negative is upwards
  ground_friction: 0.8  # Default: 0.8 | [0.1, 1.0] Friction force
  air_resistance: 0.95  # Default: 0.95 | [0.1, 1.0] Air resistance force
  max_fall_speed: 8.0   # Default: 8.0 | [0.1, 10.0] Maximum fall speed

6pt

References↩︎