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. While CNNs have been default models for encoding image observations so far, modifying the encoder architecture presents challenges, particularly due to the the necessity of identifying a new set of hyper-parameters that align with each modification. To address this problem, we propose a powerful representation learning technique for visual reinforcement learning using 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.The FNO encoder with the Efficient Rainbow algorithm achieves a median Human Normalized Score (HNS) of $26.1\%$ on the Atari100k benchmark across 26 environments, delivering a $10$-point enhancement over the CNN-based Efficient Rainbow algorithm. In the context of offline reinforcement learning Atari games, we achieve a remarkable $2.89\times$ improvement compared to sate-of-the-art transformer based models. Additionally, upon using our FNO encoder with the A2C algorithm on the ViZDoom environment, we achieve $\sim38\%$ improvement in rewards in the first $200$ episodes. Further, we match the vanilla A2C performance after just $\sim100$ episodes. We also achieve $81\%$ mean normalized score in the CARLA Autonomous Driving task (from just image sensor inputs), which is a $20$ points improvement in the absolute scale over the CNN-based PPO algorithm while requiring only $\sim55\%$ samples to match the CNN-PPO performance. We currently hold the 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 hyper-parameter tuning.