Using Semantic Segmentation to Boost Reinforcement Learning Performance

The goal of this work was to test whether semantic segmentation could act as a useful visual abstraction for reinforcement learning agents. Instead of training directly on RGB frames, I built a pipeline that converts each game frame into a semantic map and uses that representation as the agent input. This let me measure whether removing texture-level detail could improve sample efficiency and generalization across game levels.

The core hypothesis was simple: for many game-control tasks, the agent does not need texture-level detail. It needs to know which objects are present, where they are, and how they relate spatially. Semantic segmentation provides that abstraction while preserving the structure needed for control.

To train the segmentation model, I generated synthetic frames from procedurally combined cutouts, with composition rules adapted to each object type. I trained a DeepLabV3 model with a ResNet-50 backbone and used it as a preprocessing stage before a Double Deep Q-Learning agent.

The segmented representation reduced the number of episodes required for convergence and improved the final game score compared with the RGB baseline.

The colored area represents the 95% confidence interval, and the line represents the median score across five training runs.

The strongest result came from the multi-level setting. I trained the agent in a round-robin setup, where each episode sampled one of three levels. The RGB baseline failed to learn stable behavior across the three levels. The segmentation-based agent reached higher performance on all three levels and showed more consistent behavior across environments. Horizontal lines represent mean performance.

The paper also evaluates additional RL algorithms and fine-tuning experiments. The full paper is available on Springer, and the code is available on GitHub.