Mode-Dependent Rectification for Stable PPO Training

TL;DR

This paper identifies how mode-dependent layers like Batch Normalization destabilize PPO training and introduces Mode-Dependent Rectification (MDR), a simple method that stabilizes training and improves performance without changing the network architecture.

Contribution

The paper proposes MDR, a novel dual-phase training method that stabilizes PPO with mode-dependent layers, addressing a key source of instability in reinforcement learning.

Findings

MDR improves PPO stability across various tasks.

MDR extends to other mode-dependent layers.

Enhanced performance in real-world and simulated environments.

Abstract

Mode-dependent architectural components (layers that behave differently during training and evaluation, such as Batch Normalization or dropout) are commonly used in visual reinforcement learning but can destabilize on-policy optimization. We show that in Proximal Policy Optimization (PPO), discrepancies between training and evaluation behavior induced by Batch Normalization lead to policy mismatch, distributional drift, and reward collapse. We propose Mode-Dependent Rectification (MDR), a lightweight dual-phase training procedure that stabilizes PPO under mode-dependent layers without architectural changes. Experiments across procedurally generated games and real-world patch-localization tasks demonstrate that MDR consistently improves stability and performance, and extends naturally to other mode-dependent layers.