KIPPO: Koopman-Inspired Proximal Policy Optimization

TL;DR

KIPPO reduces gradient variance in policy-gradient learning for nonlinear control by learning an approximately linear latent representation via a Koopman-inspired auxiliary network. The approach decouples representation learning from the core PPO optimization and uses reconstruction, latent-space prediction, and state-space prediction losses to enforce informative, locally linear, and predictive latent dynamics. Empirical results across MuJoCo and Box2D tasks show 6–60% improvements in mean returns and 26–91% reductions in variance compared to baselines, with modest training-time overhead. This method provides a practical pathway to more stable and scalable on-policy learning in complex, non-linear environments and offers avenues for extension to off-policy and discrete domains.

Abstract

Reinforcement Learning (RL) has made significant strides in various domains, and policy gradient methods like Proximal Policy Optimization (PPO) have gained popularity due to their balance in performance, training stability, and computational efficiency. These methods directly optimize policies through gradient-based updates. However, developing effective control policies for environments with complex and non-linear dynamics remains a challenge. High variance in gradient estimates and non-convex optimization landscapes often lead to unstable learning trajectories. Koopman Operator Theory has emerged as a powerful framework for studying non-linear systems through an infinite-dimensional linear operator that acts on a higher-dimensional space of measurement functions. In contrast with their non-linear counterparts, linear systems are simpler, more predictable, and easier to analyze. In this paper, we present Koopman-Inspired Proximal Policy Optimization (KIPPO), which learns an approximately linear latent-space representation of the underlying system's dynamics while retaining essential features for effective policy learning. This is achieved through a Koopman-approximation auxiliary network that can be added to the baseline policy optimization algorithms without altering the architecture of the core policy or value function. Extensive experimental results demonstrate consistent improvements over the PPO baseline with 6-60% increased performance while reducing variability by up to 91% when evaluated on various continuous control tasks.

Figures (15)

• Figure 1: Visualization of *KIPPO's improvements relative to the *PPO baseline in terms of average performance (mean, higher is better --- left) and consistency (std., lower is better --- right) across four trials per environment.
• Figure 2: The KIPPO framework architecture. The state autoencoder (encoder $$ and decoder $$) learns a compact latent representation of environment states. The action encoder $$ maps actions to this feature space. Within the latent space, dynamics are governed by the linear state-transition matrix $$ and control matrix $$. The policy optimization algorithm operates on the encoded states $_t = (_t)$. This architecture enables the reformulation of nonlinear environments into a structure aligned with Koopman control theory Eq. \ref{['eq:koopman_control']}.
• Figure 3: The mean percent improvement of environments with various levels of complexity in performance gain (Left) and variance reduction (Right) of final returns by KIPPO compared to PPO.
• Figure 4: Hyperparameter importance scores derived from a random forest regressor.
• Figure B.1: Visualizations of the six continuous control environments used for evaluation.