One Encoder to Rule Them All: Representation Learning for Model-Free Visual Reinforcement Learning Using Fourier Neural Operators

Representation learning lies at the core of deep reinforcement learning. Although CNNs have traditionally served as the primary models for encoding image observations, modifying the encoder architecture introduces challenges, especially due to the necessity of determining a new set of hyperparameters. To address this problem, we propose a powerful representation learning technique for visual reinforcement learning utilizing Fourier Neural Operators (FNO). Our findings demonstrate that the proposed FNO encoder effectively learns representations from images that encapsulate the underlying differential equations (PDEs) governing the dynamics of the environment in an online model-free RL framework. We demonstrate the applicability of our proposed architecture by replacing the CNN image encoder in PPO, A2C, and Rainbow (a Policy Gradient, Actor-Critic, and Q-Learning RL algorithm, respectively). We achieve state-of-the-art scores (in the model-free RL setting) at both the CARLA Autonomous Driving (from image observations) benchmark and the Atari 100k benchmark. Our proposed FNO encoder is compatible with all model-free reinforcement learning algorithms, enhances both rewards and sample efficiency by implicitly learning the underlying dynamics of the environment, and eliminates the need for additional hyperparameter tuning.