Latent Space Perturbations Improve Visual Generalization in Reinforcement Learning

Visual generalization in reinforcement learning (RL) means keeping the same task performance when the visuals change, for example the background, lighting, or object appearance, while the dynamics and rewards stay the same. Humans handle such changes well, but RL agents often lose a lot of performance under these visual shifts. Inspired by efficient coding theory from neuroscience, we study how visual shifts change the encoder latent space. We find that encoders trained in fixed visual settings learn a mostly low-dimensional latent space: out-of-distribution (OOD) visuals mainly cause mean and variance shifts along the top principal components of the training latent space, instead of creating new directions. We validate this on our diagnostic environment (LSP-Env) with controlled tests for background changes, distractors, and agent appearance changes, and we also validate this on Procgen. Based on these results, we propose Latent Space Perturbation (LSP), a PCA-guided method that adds variance-scaled Gaussian noise along dominant latent directions during training. Across LSP-Env, Procgen, and RL-ViGen, LSP improves OOD performance without hurting in-distribution performance, and it does so without changing the encoder or using pixel-level augmentation.