RLax
Library of JAX-based building blocks for reinforcement learning agents
RLax (pronounced “relax”) is a JAX-based library developed by Google DeepMind that provides reusable mathematical building blocks for constructing reinforcement learning (RL) agents. Rather than implementing full algorithms, RLax focuses on the core functional operations that underpin RL methods—such as computing value functions, returns, policy gradients, and loss terms—allowing researchers to flexibly assemble their own agents. It supports both on-policy and off-policy learning, as well...