Balancing Expressivity and Robustness: Constrained Rational Activations for Reinforcement Learning

TL;DR

This paper investigates the use of trainable rational activation functions in reinforcement and continual learning, identifying a trade-off between their expressivity and stability, and proposes a constrained variant to improve robustness.

Contribution

It introduces a constrained rational activation function that balances expressivity and stability, enhancing training robustness in reinforcement and continual learning tasks.

Findings

Constrained rational activations improve training stability in continuous control environments.

Flexibility of rational activations can cause instability and feature collapse.

Different constraints influence the trade-off between expressivity and long-term retention.

Abstract

Trainable activation functions, whose parameters are optimized alongside network weights, offer increased expressivity compared to fixed activation functions. Specifically, trainable activation functions defined as ratios of polynomials (rational functions) have been proposed to enhance plasticity in reinforcement learning. However, their impact on training stability remains unclear. In this work, we study trainable rational activations in both reinforcement and continual learning settings. We find that while their flexibility enhances adaptability, it can also introduce instability, leading to overestimation in RL and feature collapse in longer continual learning scenarios. Our main result is demonstrating a trade-off between expressivity and plasticity in rational activations. To address this, we propose a constrained variant that structurally limits excessive output scaling while preserving adaptability. Experiments across MetaWorld and DeepMind Control Suite (DMC) environments show that our approach improves training stability and performance. In continual learning benchmarks, including MNIST with reshuffled labels and Split CIFAR-100, we reveal how different constraints affect the balance between expressivity and long-term retention. While preliminary experiments in discrete action domains (e.g., Atari) did not show similar instability, this suggests that the trade-off is particularly relevant for continuous control. Together, our findings provide actionable design principles for robust and adaptable trainable activations in dynamic, non-stationary environments. Code available at: https://github.com/special114/rl_rational_plasticity.