How luchris429/purejaxrl works
1. Repository purpose
- PureJaxRL is a collection of end-to-end reinforcement learning training implementations written in JAX.
- Its main architectural choice is that the environment interaction, rollout collection, and learning loop all stay in JAX, enabling jit, vmap, pmap, and scan over the full pipeline.
2. Top-level shape
- This is not a modular framework or library; it is a set of self-contained training scripts.
- The core algorithm entrypoints live under purejaxrl/, while examples/ contains notebook-based usage demos and docs/ holds performance visualizations.
3. Core algorithm subsystem
- purejaxrl/ppo.py is the central PPO implementation and likely the reference for the project’s training structure.
- purejaxrl/dqn.py provides a value-based alternative.
- purejaxrl/dpo_continuous_action.py implements the project’s discovered-policy-optimization variant for continuous actions.
- These files represent separate algorithm pipelines rather than shared class hierarchies.
4. Specialized PPO variants
- purejaxrl/ppo_continuous_action.py adapts PPO to continuous control.
- purejaxrl/ppo_minigrid.py is tailored to Minigrid-style environments.
- purejaxrl/ppo_rnn.py adds recurrence for partially observable settings or sequence modeling.
- These are parallel training workflows built around the same JAX-first design, but with environment- and policy-specific differences.
5. Environment and observation processing boundary
- purejaxrl/wrappers.py is the main preprocessing/boundary layer between raw environments and the learning code.
- This is where environment outputs are typically normalized, transformed, or adapted so the algorithms can consume them consistently.
- The repository’s architecture depends on this layer to keep the training scripts simple while still supporting multiple environment families.
6. Experimental recurrent-model subsystem
- purejaxrl/experimental/s5/ is an isolated experimental track for S5-based sequence modeling.
- s5.py contains the model implementation, ppo_s5.py integrates it into PPO, and wrappers.py provides any environment adaptation needed for that experiment.
- This subtree is intentionally separated from the main algorithm files, signaling that it is research-oriented and not yet part of the stable core.
7. Training data flow
- The typical flow is: initialize environment state → collect batched rollouts in JAX → compute returns/advantages or targets → apply policy/value updates → repeat.
- Because the whole loop remains on accelerator-friendly arrays, the system avoids Python-side stepping during the hot path.
- The architecture is built for many parallel agents/seeds, not a single online agent.
8. Vectorization and parallelism boundary
- A major architectural feature is the ability to scale over agents and seeds using JAX transformations rather than distributed service infrastructure.
- This is important to the repo’s design: performance comes from data-parallel execution within one training program, not from external orchestration.
- The README’s performance claims and plots reflect this vectorized training model.
9. Usage and entry surfaces
- examples/walkthrough.ipynb is the primary onboarding path and likely shows how to run and modify a training pipeline.
- examples/brax_minatar.ipynb demonstrates application to specific environments from the JAX RL ecosystem, especially Brax and MinAtar.
- These notebooks are the practical “front door” for the repository’s architecture.
10. Documentation and ecosystem context
- RESOURCES.md maps the repository into the broader JAX RL ecosystem.
- docs/ contains benchmark figures that document performance characteristics, which are a first-class part of the project’s story.
- requirements.txt and the README indicate the key runtime dependency is JAX, with environment-specific packages layered on top.
11. Architectural summary for diagramming
- A good diagram would show:
- notebook/examples at the edge,
- a set of standalone algorithm scripts in the middle,
- wrappers as the environment adaptation boundary,
- and JAX-based vectorized environment/rollout/update execution as the core engine.
- The experimental S5 subtree should be drawn as a separate branch off the main PPO line, not as part of the primary stable path.
Main entry points: wrappers.py Wrappers · ppo_s5.py PPO S5 · ppo.py PPO · ppo_continuous_action.py PPO cont.
Components 17
Entry points
Primary onboarding and usage surfaces.
- Walkthrough · Main onboarding notebook for running and modifying training loops.
examples/walkthrough.ipynb - Brax/MinAtar · Application-focused notebook showing environment-specific usage.
examples/brax_minatar.ipynb
Core training
Main JAX-first training pipelines and shared boundary layers.
- Wrappers · Shared preprocessing and environment adaptation layer.
purejaxrl/wrappers.py - JAX loop · Conceptual core: batched rollout collection, targets, and updates stay inside JAX transformations.
- PPO · Reference PPO training pipeline and central project structure.
purejaxrl/ppo.py - DQN · Alternative training pipeline using value-based updates.
purejaxrl/dqn.py - PPO cont. · PPO variant adapted for continuous action spaces.
purejaxrl/ppo_continuous_action.py - PPO Minigrid · PPO workflow tailored to Minigrid-style environments.
purejaxrl/ppo_minigrid.py - PPO RNN · PPO variant with recurrence for partial observability and sequence modeling.
purejaxrl/ppo_rnn.py - DPO cont. · Discovered-policy-optimization variant for continuous actions.
purejaxrl/dpo_continuous_action.py
Experimental S5
Research branch for sequence-model experiments.
- S5 model · Experimental S5-based sequence model implementation.
purejaxrl/experimental/s5/s5.py - PPO S5 · PPO integration that uses the S5 model.
purejaxrl/experimental/s5/ppo_s5.py - S5 wrappers · Environment adaptation layer for the S5 experiment.
purejaxrl/experimental/s5/wrappers.py - S5 notes · Documentation for the experimental S5 branch.
purejaxrl/experimental/s5/README.md
Docs & context
Repository documentation and ecosystem references.
- README · Top-level explanation of the project, training model, and performance story.
README.md - Resources · Links the repository into the broader JAX RL ecosystem.
RESOURCES.md - Benchmarks · Performance plots and benchmark visualizations.
docs
Connections 21
- Walkthrough → PPO: explains · Primary notebook demonstrates the reference PPO flow.
- Brax/MinAtar → PPO cont.: demonstrates · Shows application to specific environment families.
- README → PPO: summarizes · Introduces the main training architecture and design choices.
- Resources → README: context · Places the repo in the wider ecosystem.
- Benchmarks → README: evidence · Benchmarks support the performance claims in the README.
- Wrappers → PPO: feeds · Provides normalized environment outputs for training.
- Wrappers → PPO cont.: feeds · Shared environment boundary for specialized PPO variants.
- Wrappers → PPO Minigrid: feeds · Adapts Minigrid-like observations and actions.
- Wrappers → PPO RNN: feeds · Supplies sequence-ready environment interaction.
- Wrappers → DPO cont.: feeds · Shared preprocessing for the DPO pipeline.
- PPO → JAX loop: uses · Core PPO keeps rollout and update logic inside JAX.
- DQN → JAX loop: uses · Value-based training also runs in the vectorized JAX loop.
- PPO cont. → JAX loop: uses · Continuous-control PPO remains accelerator-friendly.
- PPO Minigrid → JAX loop: uses · Task-specific PPO retains the same batched execution model.
- PPO RNN → JAX loop: uses · Recurrent PPO still executes rollout and learning in JAX.
- DPO cont. → JAX loop: uses · DPO variant follows the same end-to-end training pattern.
- S5 model → PPO S5: powers · S5 sequence model is embedded into the PPO experiment.
- S5 wrappers → PPO S5: feeds · Experimental boundary layer for the S5 pipeline.
- PPO S5 → PPO: branches from · Experimental track extends the core PPO structure.
- S5 notes → PPO S5: documents · Experimental notes explain the S5 workflow.
- PPO S5 → JAX loop: uses · Experimental PPO still relies on the same JAX-first execution model.