KestRL (Kestrel)
JAX/Flax Reinforcement Learning: clean implementations, hardware-accelerated training.

KestRL is a JAX/Flax reinforcement learning library built around one idea: gradient updates should never wait on Python. It is designed for clean algorithm implementations, reproducible research, and easy extension.

Networks are defined with Flax NNX — full OOP, readable, composable. Before training starts, the entire model bundle is split into a static graph definition and a mutable state PyTree exactly once. The graph definition is captured in a closure; only the state flows in and out of JIT. The result is zero retracing, zero Python overhead per gradient step.

For Brax and MJX environments, this extends further: the rollout, buffer write, gradient update, and soft target update all fold into a single lax.scan kernel. Multi-seed training becomes jax.vmap over independent carries — multiple seeds run in one dispatch, with one host sync per epoch for logging.

Wall-clock time for 1 M environment steps — Standard Route vs eager baselines (left) and Compiled Route vs functional frameworks (right):

• Standard Route outperforms CleanRL and SB3 on CPU-stepped Gymnasium environments.
• Compiled Route outperforms Rejax and matches Stoix on Brax across all seed counts.
• Scaling from N=1 to N=100 parallel seeds costs only ~40 extra minutes total: 80–88× amortized speedup over sequential baselines.

Learning curves — algorithmic integrity across seeds:

KestRL Standard tracks CleanRL and SB3 precisely; KestRL Compiled mirrors Rejax and Stoix. Multi-seed vectorization introduces no algorithmic drift.

⚠️ Requires a JAX-native environment (Brax, MJX, Gymnax, Jumanji, …).
~ Partial support.

KestRL is the only JAX library that supports standard Gymnasium environments with JIT-compiled updates and full on-device compilation with multi-seed vmap.

Standard Gymnasium environments, NumPy replay buffer, JIT-compiled gradient updates. The env.step crosses Python; everything else stays on-device.

agent = SAC(env, algo_cfg, seed=0)
trainer = Trainer(agent, config, writer=writer)
trainer.train()

Entire training loop on-device — rollout, buffer write, gradient update, soft target update — all in one lax.scan kernel.

trainer = CompiledTrainer(env, config, writer=writer)
trainer.train(seed=0)                           # single seed, live progress
trainer.train_seeds_live(num_seeds=100, seed=0) # 100 seeds, vmapped epochs

Under the hood, both routes share the same Frozen-Bundle Pattern:

# define with NNX — readable, standard OOP
self.actor  = MultiHeadMLP(obs_dim, head_configs, hidden_dims, rngs=rngs)
self.critic = MLP(critic_in, 1, hidden_dims, rngs=rngs)

# split exactly once before training
bundle_gd, bundle_state = nnx.split(bundle)

# compile a pure update — bundle_gd is captured, never re-traced
self._jit_update = _make_frozen_update(bundle_gd, ...)

# each step: only flat state crosses the JIT boundary
bundle_state, metrics = self._jit_update(bundle_state, batch, key)

In the Compiled Route, bundle_state becomes a field in the lax.scan carry. Same pattern, wider scope.

KestRL ships a full implementation of PB-SAC (Posterior Bootstrap SAC), a PAC-Bayes actor-critic method. PB-SAC maintains a block-diagonal low-rank posterior over the actor weights, optimises a PAC-Bayes bound to certify policy performance, and uses posterior-guided exploration during a critic-recalibration phase. Both Standard and Compiled routes are supported.

KestRL was the experimental platform for the paper introducing PB-SAC, accepted at ICML 2026: arxiv.org/abs/2510.10544.

git clone https://github.com/adidi24/KestRL.git
cd KestRL
uv sync

Python 3.10–3.12. Ships with the JAX CPU backend; swap in jax[cuda12] or jax[tpu] for accelerators.

# SAC on HalfCheetah-v5
uv run python -m kestrl.experiments.run

# SAC on CartPole-v1 (discrete)
uv run python -m kestrl.experiments.run algorithm=sac_cartpole environment=cartpole

# Compiled SAC on Brax Ant, 4 seeds
uv run python -m kestrl.experiments.run environment=brax algorithm=sac_mujoco experiment.num_seeds=4

# W&B tracking
uv run python -m kestrl.experiments.run experiment.track=true

Configuration via Hydra, config groups under src/kestrl/configs/.

The Frozen-Bundle Pattern eliminates @nnx.jit graph-traversal overhead by splitting the algorithm state once before training. The GraphDef is captured in compiled closures; only the flat state PyTree crosses the JIT boundary at each step.

Single-step boundary overhead (H100, synchronous):

Loop scaling — total time for K steps (ms):

nnx.cached_partial is marginally faster per isolated step but mutates state via Python-side references, breaking lax.scan and jax.vmap compatibility. The Frozen-Bundle drops 1 000-step loop overhead from 2.4 s to 86 ms — a 28× reduction.

JAX

• rejax — same compiled-training idea, broader algorithm coverage, requires JAX-native environments
• PureJaxRL — where the end-to-end JAX training idea started, single-file and easy to read
• Stoix — distributed actor-learner at scale, TPU/GPU clusters

PyTorch

• CleanRL — single-file implementations, the best place to read how an algorithm actually works
• Stable-Baselines3 — the production standard
• ObjectRL — object-oriented PyTorch RL library with a focus on modularity and clean architecture

KestRL is currently solo-built and maintained. Contributions and feedback are welcome.

No public community is available yet. If you're interested in joining the private Discord server, reach out on X (@zitouniabdelkr1) or LinkedIn (Abdelkrim ZITOUNI).

Planned:

• More algorithms — PPO, DQN
• Documentation — full API reference and tutorials
• Contributing guide — structured onboarding for contributors

Apache 2.0