# luchris429/purejaxrl architecture

> Architecture diagram of the GitHub repository luchris429/purejaxrl, generated by GitDiagram from its file tree, README and a sample of source files (last updated 2026-06-22). It is an AI-made overview: check details against the source.

- Interactive diagram: https://gitdiagram.com/luchris429/purejaxrl
- Repository: https://github.com/luchris429/purejaxrl

## Overview

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: [`purejaxrl/wrappers.py`](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/wrappers.py) (Wrappers), [`purejaxrl/experimental/s5/ppo_s5.py`](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/ppo_s5.py) (PPO S5), [`purejaxrl/ppo.py`](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo.py) (PPO), [`purejaxrl/ppo_continuous_action.py`](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_continuous_action.py) (PPO cont.)

## Diagram (Mermaid)

```mermaid
flowchart TD

subgraph group_frontdoor["Entry points"]
  node_examples_walkthrough["Walkthrough<br/>notebook demo<br/>[walkthrough.ipynb]"]
  node_examples_brax_minatar["Brax/MinAtar<br/>notebook demo<br/>[brax_minatar.ipynb]"]
end

subgraph group_core["Core training"]
  node_wrappers["Wrappers<br/>env boundary<br/>[wrappers.py]"]
  node_jax_loop(("JAX loop<br/>vectorized engine"))
  node_ppo_core["PPO<br/>main algorithm<br/>[ppo.py]"]
  node_dqn["DQN<br/>value-based RL<br/>[dqn.py]"]
  node_ppo_cont["PPO cont.<br/>continuous PPO"]
  node_ppo_minigrid["PPO Minigrid<br/>task-specific PPO<br/>[ppo_minigrid.py]"]
  node_ppo_rnn["PPO RNN<br/>recurrent PPO<br/>[ppo_rnn.py]"]
  node_dpo_cont["DPO cont.<br/>continuous DPO"]
end

subgraph group_experimental["Experimental S5"]
  node_s5_model["S5 model<br/>sequence model<br/>[s5.py]"]
  node_s5_ppo["PPO S5<br/>experimental PPO<br/>[ppo_s5.py]"]
  node_s5_wrappers["S5 wrappers<br/>exp env boundary<br/>[wrappers.py]"]
  node_s5_readme["S5 notes<br/>experiment docs<br/>[README.md]"]
end

subgraph group_context["Docs & context"]
  node_readme["README<br/>project overview<br/>[README.md]"]
  node_resources["Resources<br/>ecosystem map<br/>[RESOURCES.md]"]
  node_docs["Benchmarks<br/>figures"]
end

node_examples_walkthrough -->|"explains"| node_ppo_core
node_examples_brax_minatar -->|"demonstrates"| node_ppo_cont
node_readme -->|"summarizes"| node_ppo_core
node_resources -->|"context"| node_readme
node_docs -.->|"evidence"| node_readme
node_wrappers -->|"feeds"| node_ppo_core
node_wrappers -->|"feeds"| node_ppo_cont
node_wrappers -->|"feeds"| node_ppo_minigrid
node_wrappers -->|"feeds"| node_ppo_rnn
node_wrappers -->|"feeds"| node_dpo_cont
node_ppo_core -->|"uses"| node_jax_loop
node_dqn -->|"uses"| node_jax_loop
node_ppo_cont -->|"uses"| node_jax_loop
node_ppo_minigrid -->|"uses"| node_jax_loop
node_ppo_rnn -->|"uses"| node_jax_loop
node_dpo_cont -->|"uses"| node_jax_loop
node_s5_model -->|"powers"| node_s5_ppo
node_s5_wrappers -->|"feeds"| node_s5_ppo
node_s5_ppo -.->|"branches from"| node_ppo_core
node_s5_readme -->|"documents"| node_s5_ppo
node_s5_ppo -->|"uses"| node_jax_loop

click node_examples_walkthrough "https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb"
click node_examples_brax_minatar "https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb"
click node_readme "https://github.com/luchris429/purejaxrl/blob/main/README.md"
click node_resources "https://github.com/luchris429/purejaxrl/blob/main/RESOURCES.md"
click node_docs "https://github.com/luchris429/purejaxrl/tree/main/docs"
click node_wrappers "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/wrappers.py"
click node_ppo_core "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo.py"
click node_dqn "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/dqn.py"
click node_ppo_cont "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_continuous_action.py"
click node_ppo_minigrid "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_minigrid.py"
click node_ppo_rnn "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_rnn.py"
click node_dpo_cont "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/dpo_continuous_action.py"
click node_s5_model "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/s5.py"
click node_s5_ppo "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/ppo_s5.py"
click node_s5_wrappers "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/wrappers.py"
click node_s5_readme "https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/README.md"

classDef toneNeutral fill:#f8fafc,stroke:#334155,stroke-width:1.5px,color:#0f172a
classDef toneBlue fill:#dbeafe,stroke:#2563eb,stroke-width:1.5px,color:#172554
classDef toneAmber fill:#fef3c7,stroke:#d97706,stroke-width:1.5px,color:#78350f
classDef toneMint fill:#dcfce7,stroke:#16a34a,stroke-width:1.5px,color:#14532d
classDef toneRose fill:#ffe4e6,stroke:#e11d48,stroke-width:1.5px,color:#881337
classDef toneIndigo fill:#e0e7ff,stroke:#4f46e5,stroke-width:1.5px,color:#312e81
classDef toneTeal fill:#ccfbf1,stroke:#0f766e,stroke-width:1.5px,color:#134e4a
class node_examples_walkthrough,node_examples_brax_minatar toneBlue
class node_wrappers,node_jax_loop,node_ppo_core,node_dqn,node_ppo_cont,node_ppo_minigrid,node_ppo_rnn,node_dpo_cont toneAmber
class node_s5_model,node_s5_ppo,node_s5_wrappers,node_s5_readme toneMint
class node_readme,node_resources,node_docs toneRose
```

## Components

### Entry points

Primary onboarding and usage surfaces.

- [Walkthrough](https://github.com/luchris429/purejaxrl/blob/main/examples/walkthrough.ipynb): notebook demo · Main onboarding notebook for running and modifying training loops. · `examples/walkthrough.ipynb`
- [Brax/MinAtar](https://github.com/luchris429/purejaxrl/blob/main/examples/brax_minatar.ipynb): notebook demo · Application-focused notebook showing environment-specific usage. · `examples/brax_minatar.ipynb`

### Core training

Main JAX-first training pipelines and shared boundary layers.

- [Wrappers](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/wrappers.py): env boundary · Shared preprocessing and environment adaptation layer. · `purejaxrl/wrappers.py`
- **JAX loop**: vectorized engine · Conceptual core: batched rollout collection, targets, and updates stay inside JAX transformations.
- [PPO](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo.py): main algorithm · Reference PPO training pipeline and central project structure. · `purejaxrl/ppo.py`
- [DQN](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/dqn.py): value-based RL · Alternative training pipeline using value-based updates. · `purejaxrl/dqn.py`
- [PPO cont.](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_continuous_action.py): continuous PPO · PPO variant adapted for continuous action spaces. · `purejaxrl/ppo_continuous_action.py`
- [PPO Minigrid](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_minigrid.py): task-specific PPO · PPO workflow tailored to Minigrid-style environments. · `purejaxrl/ppo_minigrid.py`
- [PPO RNN](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/ppo_rnn.py): recurrent PPO · PPO variant with recurrence for partial observability and sequence modeling. · `purejaxrl/ppo_rnn.py`
- [DPO cont.](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/dpo_continuous_action.py): continuous DPO · Discovered-policy-optimization variant for continuous actions. · `purejaxrl/dpo_continuous_action.py`

### Experimental S5

Research branch for sequence-model experiments.

- [S5 model](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/s5.py): sequence model · Experimental S5-based sequence model implementation. · `purejaxrl/experimental/s5/s5.py`
- [PPO S5](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/ppo_s5.py): experimental PPO · PPO integration that uses the S5 model. · `purejaxrl/experimental/s5/ppo_s5.py`
- [S5 wrappers](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/wrappers.py): exp env boundary · Environment adaptation layer for the S5 experiment. · `purejaxrl/experimental/s5/wrappers.py`
- [S5 notes](https://github.com/luchris429/purejaxrl/blob/main/purejaxrl/experimental/s5/README.md): experiment docs · Documentation for the experimental S5 branch. · `purejaxrl/experimental/s5/README.md`

### Docs & context

Repository documentation and ecosystem references.

- [README](https://github.com/luchris429/purejaxrl/blob/main/README.md): project overview · Top-level explanation of the project, training model, and performance story. · `README.md`
- [Resources](https://github.com/luchris429/purejaxrl/blob/main/RESOURCES.md): ecosystem map · Links the repository into the broader JAX RL ecosystem. · `RESOURCES.md`
- [Benchmarks](https://github.com/luchris429/purejaxrl/tree/main/docs): figures · Performance plots and benchmark visualizations. · `docs`

## Connections

- 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.