Abstract
Training reinforcement learning (RL) agents often requires significantcomputational resources and extended training times. To address this, we buildupon the foundation laid by Google Brain's Sensory Neuron, which introduced anovel neural architecture for reinforcement learning tasks that maintainedpermutation in-variance in the sensory neuron system. While the baseline modeldemonstrated significant performance improvements over traditional approaches,we identified opportunities to enhance the efficiency of the learning processfurther. We propose a modified attention mechanism incorporating a non-lineartransformation of the key vectors (K) using a mapping function, resulting in anew set of key vectors (K'). This non-linear mapping enhances therepresentational capacity of the attention mechanism, allowing the model toencode more complex feature interactions and accelerating convergence withoutcompromising performance. Our enhanced model demonstrates significantimprovements in learning efficiency, showcasing the potential for non-linearattention mechanisms in advancing reinforcement learning algorithms.