Interpreting Reinforcement Learning Model Behavior via Koopman with Control

Abstract

Reinforcement learning (RL) models have shown the capability of learning complex behaviors, but quantitatively assessing those behaviors - which is critical for safety assurance and the discovery of novel strategies - is challenging. By viewing RL models as control systems, we hypothesize that data-driven approximations of their associated Koopman operators may provide dynamical information about their behavior, thus enabling greater interpretability. To test this, we apply the Koopman with control framework to RL models trained on several standard benchmark environments and demonstrate that properties of the fit linear control models, such as stability and controllability, evolve during training in a task dependent manner. Comparing these metrics across different training epochs or across differently optimized RL models enables an understanding of how they differ. In addition, we find cases where - even when the reward achieved by the RL model is static - the stability and controllability is nonetheless evolving, predicting increased reward with further training. This suggests that these metrics may be able to serve as hidden progress measures, a core idea in mechanistic interpretability. Taken together, our results illustrate that the Koopman with control framework provides a comprehensive way in which to analyze and interpret the behavior of RL models, particularly across training.

Figures (5)

• Figure 1: Schematic illustration of how different behaviors of RL models may have different associated Koopman eigenvalues. (A) Three example dynamical behaviors ("falling", "hovering", "landing") that are possible in the LunarLander RL environment. First column is at the start of the trial ($T = 0$) and second column is $20$ time steps later ($T = 20$). (B) Possible Koopman eigenvalues associated with each of the behaviors in (A).
• Figure 2: Performant RL models exhibit more stable trajectories, with greater controllability, across training on CartPole. (A) Schematic illustration of the CartPole environment. (B) Median reward, as a function of training epoch, for RL models optimized using PPO (blue) and A2C (red). (C) Maximum eigenvalue norm, as a function of epoch. (D) Normalized rank of controllability matrix, as a function of epoch. (B)--(D) Solid line represents mean and shaded area represents $\pm$ standard error across 25 independently trained RL models.
• Figure 3: Koopman eigenvalue norm and controllability matrix rank contain information that can act as a hidden progress measure. (A) Median reward, as a function of training epoch, for RL models optimized using A2C. Orange shaded area denotes training time window where little change is seen in the reward. (B) Maximum eigenvalue norm and normalized controllability matrix rank, as a function of epoch. Orange shaded area denotes same training time window as in (A). Note that, for these two metrics, there is an increase in stability and controllability, during this training time window.
• Figure 4: Performant RL models exhibit less stable trajectories, with a greater controllability, across training on Acrobot. (A) Schematic illustration of the Acrobot environment. (B) Reward, as a function of training epoch, for RL models optimized using PPO (blue) and A2C (red). (C) Maximum norm of eigenvalues, as a function of training epoch. (D) Normalized rank of controllability matrix, as a function of epoch. (B)--(D) Solid line represents mean and shaded area represents $\pm$ standard error across 25 independently trained RL models.
• Figure 5: Performant RL models exhibit more stable trajectories, with greater controllability, across training on LunarLander. (A) Schematic illustration of the LunarLander environment. (B) Median reward, as a function of training epoch, for RL models optimized using PPO (blue) and A2C (red). (C) Maximum norm of eigenvalues, as a function of epoch. (D) Normalized ank of controllability matrix, as a function of epoch. (B)--(D) Solid line represents mean and shaded area represents $\pm$ standard error across 25 independently trained RL models.