Pixel-Based Reinforcement Learning

• Pixel-based reinforcement learning is a framework where agents learn directly from high-dimensional pixel data, enabling focused studies on visual generalization.
• KAGE-Env and KAGE-Bench decompose visual observations into approximately 90 controllable axes to isolate and evaluate specific visual shifts.
• The integration of JAX-parallelization and vectorization achieves high simulation throughput, facilitating rapid empirical testing and robust RL algorithm development.

Pixel-based reinforcement learning (RL) refers to RL frameworks in which agents perceive their environment directly through high-dimensional pixel observations, rather than structured or low-dimensional state representations. Addressing the visual generalization challenge—where agents must maintain performance across visually distinct but functionally identical scenarios—has become a central focus. KAGE-Env and KAGE-Bench exemplify state-of-the-art efforts to systematically decompose and rigorously evaluate the visual generalization capacity of pixel-based RL agents by factorizing the observation process into independently controllable visual axes, with a strictly fixed latent control problem (Cherepanov et al., 20 Jan 2026).

1. Observation Factorization and Visual POMDP Formulation

The KAGE-Env paradigm frames the environment as a family of visual partially observable Markov decision processes (POMDPs) indexed by a vector of visual parameters $\xi \in \Xi$. The core objects are:

• Latent state space $S$ and states $s \in S$ (e.g., position, velocity).
• Action space $A = \{0, \dots, 7\}$, encoding core action primitives as a bitmask.
• Observation space $$\Omega \subseteq \{0, \dots, 255\}^{H \times W \times 3}$$ with default $H = W = 128$ (RGB images).
• Visual parameter space $V = V_1 \times \cdots \times V_k$ (background, filters, lighting, sprites, distractors, layout, etc.).
• Rendering process (observation kernel) $O_\xi(\cdot|s)$ mapping from $S$ to $\Omega$, concretely implemented with a deterministic rendering function $o = f(s; \xi, \text{key})$.

Crucially, the latent transition kernel $P(s'|s,a)$ and reward function $r(s,a)$ are independent of $\xi$, ensuring that the underlying MDP remains fixed while the entire visual channel is modifiable. This allows any change in agent performance under visual shifts to be attributed unambiguously to visual generalization, not to confounding changes in dynamics or rewards.

A reactive pixel policy $\pi(a|o)$ in environment $\mathcal{M}_\xi$ induces a state-conditional policy $\pi_\xi(a|s) = \int_{\Omega} \pi(a|o) O_\xi(do|s)$ in the latent MDP. Theorem 1 of (Cherepanov et al., 20 Jan 2026) demonstrates formal equivalence at the level of $(s_t, a_t)$ trajectories and returns.

2. Visual Axes: Controllable Visual Dimensions

KAGE-Env exposes approximately 90 rendering and observation parameters as independent axes, which can be precisely manipulated. The principal axes are:

Axis	Key Modes/Controls	Example Range/Description
Background	"black"/"color"/"noise"/"image", parallax, tiling, switch frequency	Single color, uniform random, 128 photos
Agent	27 animated sprites, 9 shapes × 21 colors, animation fps, rotation	Skin, color, shape, pose
NPCs/Dist.	Sprite dirs, per-episode NPCs, sticky NPCs, distractor count/types/colors/kinematics	Moving shapes, world-fixed/sticky
Photometric	Brightness, contrast, gamma, saturation, hue, color temp, noise, blur, pixelation, pop filters	Stochastic and deterministic transforms
Lighting	Point light count (1–5), intensity, radius, falloff, color	Dynamic relighting
Layout	Level length, ground height, step size/prob, stair width, color palette	Procedural geometries
Physics	Gravity, move_speed, jump_force, friction, air resistance, max fall speed (optional)	Dynamical system

Each train–eval pair isolates a shift along a single axis, ensuring causal attribution of downstream generalization failures or successes to the manipulated visual property.

3. Level Construction, Sampling, and Reward Specification

Levels are generated procedurally as side-scrolling layouts (sequences of horizontal segments and staircases) with parameters directly controlling step size, frequency, ground height, and platform thickness. The world is "push-scrolling," with the camera following the agent's progress and exposing new geometry incrementally.

At episode reset, the visual parameter vector $\xi$ is sampled from either a deterministic $D_\text{train}$ or $D_\text{eval}$, defined by KAGE-Bench to fix all but one axis per experiment.

Each observation is rendered in real time via a custom JAX shader pipeline, producing $128 \times 128$ uint8 RGB frames per step.

The per-step reward is

$$r_t = \alpha_1 \max \left\{0, x_{t+1} - x_t^{\text{max}} \right\} - \left[\alpha_2 1_{\{\text{Jump}(a_t)\}} + \alpha_3 + \alpha_4 1_{\{\text{idling}\}} \right]$$

where $x_t$ is horizontal position, $x_t^{\text{max}}$ is maximum reached so far, and $\alpha$ are configurable penalty/reward coefficients. Episodes terminate only at a fixed time horizon.

4. JAX-Parallel Environment API and Vectorization

KAGE-Env is implemented entirely in JAX and designed to maximize high-throughput simulation and reproducibility:

• Python API: Environments are instantiated from YAML configurations and support reset(key) and step(state, action) methods. Internal state is a JAX pytree, ensuring functional purity.
• Vectorization: Using JAX primitives (jit, vmap, lax.scan), bulk operations over up to $2^{16}$ parallel environments can be fused for maximal GPU utilization, avoiding Python overhead.
• Rendering: Observation generation (rendering) is a pure JAX function, thus offline video batch rendering via jit(vmap(env.render)) is direct.
• Performance: On an NVIDIA H100 (80 GB), throughput reaches up to 33M environment steps/sec with simple visual settings, and 10–15M steps/sec with maximal axis complexity and $2^{16}$ environments. On Apple M3 Pro CPU with $\sim$1024 envs: 1–2M steps/sec.

Such throughput enables comprehensive studies: 10 seeds × 34 axis configuration pairs × 25M steps per run executed in a few GPU-hours.

5. Benchmarking Visual Generalization: KAGE-Bench Protocol

KAGE-Bench comprises six “known-axis” benchmark suites, each with 34 train–eval configuration pairs corresponding to isolated visual axis modifications. In each suite, only a single visual property is changed at evaluation, allowing for sharp tests of axis-dependent generalization.

Empirical results with a standard PPO-CNN baseline indicate pronounced axis-dependence: background and photometric shifts commonly induce failure (collapsing success rate), while agent-appearance shifts are more benign. Some visual shifts break task completion while preserving forward motion, demonstrating that cumulative return alone may obscure subtle failures of generalization.

A plausible implication is that fine-grained or axis-specific evaluation protocols are necessary to fully characterize and stress-test the visual robustness of pixel-based RL agents (Cherepanov et al., 20 Jan 2026).

6. Integration with RL Algorithms

The KAGE-Env API supports plug-and-play integration with standard RL training loops, such as PPO. Example protocol:

from kage_bench import KAGE_Env, load_config_from_yaml config = load_config_from_yaml("train_bg_black.yaml") env = KAGE_Env(config) N_env = 256 reset = jit(vmap(env.reset)) step = jit(vmap(env.step)) keys = random.split(random.PRNGKey(seed), N_env) obs, info = reset(keys) states = info["state"] for update in range(num_updates): for t in range(T_hor): logits = policy_apply(params, obs/255.0) actions = sample_actions(logits, keys) obs2, rews, dones, truncs, info2 = step(states, actions) states = info2["state"] obs = obs2 params = ppo_update(params, buffer)

Hyperparameters adhere to common benchmarks (CleanRL/PPO-CNN: 128 envs, 128 steps, learning rate $5 \times 10^{-4}$, batch size $128 \times 128$, 3 epochs, clip 0.2). The user consults the Appendix of (Cherepanov et al., 20 Jan 2026) for precise experimental details, configuration files, and learning curves.

7. Significance and Implications for RL Research

Systematic visual generalization evaluation, as instantiated in KAGE-Bench, enables the field to decouple and rigorously attribute generalization failures to specific visual axes, free from confounding changes in underlying behavior or reward. The KAGE paradigm both accelerates empirical iteration (by orders of magnitude) and clarifies the visual robustness landscape for pixel policy learning. It provides a clean testbed for architectural, algorithmic, and augmentation proposals to close the visual generalization gap in RL, and supports both fast and reproducible sweeps over high-dimensional visual parameter spaces (Cherepanov et al., 20 Jan 2026).