Learning Switched Koopman Models for Control of Entity-Based Systems

One problem that arises in control engineering is that of controlling a system for which the dynamics are unknown. Such a problem favors a data-driven approach, such as can be done through the use of the Koopman operator. We present a switched Koopman model, applicable to systems with discrete sets of inputs, that gives rise to an optimal control problem with a piecewise affine value function. This structure provides an efficient representation and enables a heuristic pruning algorithm that avoids the exponential complexity of finding the true optimal solution. We use density-like observables that are defined through the notion of entity-based systems: systems whose state is composed of a possibly varying number of entities that can be grouped into classes of like-entities. This encompasses many systems, including arcade games, a common benchmark used in reinforcement learning. We find that our approach requires much less training than commonly used for reinforcement learning. The Koopman approach also has the advantage of being agnostic to the control objective, which allows the objective to be changed without needing to retrain the model.