Metadata-Version: 2.4
Name: atomic-rl
Version: 0.2.0
Summary: Reinforcement learning utilities for Deep RL and the Alberta Plan: readable implementations of DQN, Rainbow, PPO, and MCTS.
Author: Jonathan Lamontange-Kratz
License-Expression: MIT
Project-URL: Homepage, https://github.com/epicgamer17/atomic-rl
Project-URL: Repository, https://github.com/epicgamer17/atomic-rl
Classifier: Development Status :: 3 - Alpha
Classifier: Intended Audience :: Developers
Classifier: Intended Audience :: Science/Research
Classifier: Programming Language :: Python :: 3
Classifier: Programming Language :: Python :: 3.9
Classifier: Programming Language :: Python :: 3.10
Classifier: Programming Language :: Python :: 3.11
Classifier: Programming Language :: Python :: 3.12
Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
Requires-Python: >=3.9
Description-Content-Type: text/markdown
License-File: LICENSE
Requires-Dist: torch
Requires-Dist: numpy
Requires-Dist: gymnasium
Requires-Dist: tensordict
Requires-Dist: matplotlib
Requires-Dist: wandb
Requires-Dist: tqdm
Requires-Dist: pygame
Provides-Extra: envs
Requires-Dist: pettingzoo; extra == "envs"
Requires-Dist: catanatron; extra == "envs"
Requires-Dist: rlcard; extra == "envs"
Provides-Extra: examples
Requires-Dist: ray; extra == "examples"
Requires-Dist: torchvision; extra == "examples"
Requires-Dist: gym_microrts; extra == "examples"
Provides-Extra: test
Requires-Dist: pytest>=7.0; extra == "test"
Requires-Dist: pytest-cov; extra == "test"
Requires-Dist: pytest-rerunfailures; extra == "test"
Provides-Extra: plot
Requires-Dist: SciencePlots>=2.1.0; extra == "plot"
Requires-Dist: catppuccin-matplotlib>=0.5; extra == "plot"
Requires-Dist: mplcyberpunk>=0.7.1; extra == "plot"
Requires-Dist: matplotx>=0.3.10; extra == "plot"
Requires-Dist: aquarel>=0.0.6; extra == "plot"
Requires-Dist: qbstyles>=0.1.4; extra == "plot"
Dynamic: license-file

[![Interactive Research Labs](https://img.shields.io/badge/Live_Demos-Interactive_Research_Labs-blue?style=for-the-badge&logo=flask)](https://kratzj.vercel.app/labs)
[![Website](https://img.shields.io/badge/Website-kratzj.vercel.app-10b981?style=for-the-badge&logo=vercel)](https://kratzj.vercel.app)
[![Buy Me A Coffee](https://img.shields.io/badge/Buy_Me_A_Coffee-FFDD00?style=for-the-badge&logo=buy-me-a-coffee&logoColor=black)](https://www.buymeacoffee.com/epicgamer17)
[![Patreon](https://img.shields.io/badge/Patreon-F1465A?style=for-the-badge&logo=patreon&logoColor=white)](https://www.patreon.com/epicgamer17)

# Modular RL (Functional)

A high-performance, researcher-centric Reinforcement Learning library for PyTorch built on the **Functional Core, Imperative Shell** design pattern. Think of it as **RLax for PyTorch**—engineered to eliminate the rigidity of deep OOP frameworks and the chaos of monolithic single-file copy-pasting.

---

## 📦 Installation

```bash
pip install atomic-rl
```

Or, to run the examples and experiments directly from a clone of this repository:

```bash
git clone https://github.com/epicgamer17/modular-rl.git
cd modular-rl
pip install -e ".[envs,examples,plot,test]"
```

## 🚀 Quickstart

The library ships as three flat, top-level import packages: `atomic_rl` (pure mathematical primitives), `networks` (network building blocks), and `envs` (environments, streams, and wrappers).

```python
import torch
from atomic_rl.networks.initialization import layer_init_
from atomic_rl.utils import set_seed
from atomic_rl.action_selection import argmax_selector, with_epsilon_greedy
from atomic_rl.td import compute_v_td_target
from atomic_rl.networks.resnet import ResNetBlock

set_seed(42)

# One-step TD value targets: r + gamma * V(s') * (1 - done)
next_values = torch.tensor([3.0, 4.0])
rewards = torch.tensor([0.5, 1.0])
terminated = torch.tensor([0.0, 1.0])
gamma = torch.tensor([0.9, 0.9])
targets = compute_v_td_target(next_values, rewards, terminated, gamma)
# -> tensor([3.2, 1.0])

# Greedy action selection with epsilon-greedy wrapper over Q-values [B, A]
select = with_epsilon_greedy(argmax_selector)
q_values = torch.randn(1, 4)
actions, _ = select(predictions=q_values, epsilon=0.1, num_actions=4,
                    generator=torch.Generator().manual_seed(0))
```

Complete, runnable algorithms (DQN, PPO, A2C, AlphaZero, MuZero, TD-learning, and more) live in the [`examples/`](examples/) directory.

---

## 🌍 Supported Environments & Ecosystem

`atomic-rl` integrates seamlessly with standard Python RL ecosystem libraries rather than re-inventing the wheel:

| Ecosystem / Library | Target Domain | Status & Integration |
|---|---|---|
| **Gymnasium** | Classic Control, Box2D, Discrete & Continuous MDPs | Full Native Support (`gym.make(...)`) |
| **MuJoCo** (`gymnasium[mujoco]`) | High-Dimensional Continuous Robotics | Full Support (`HalfCheetah-v4`, `Hopper-v4`, `Ant-v4`) |
| **PettingZoo** (`pettingzoo[classic]`) | Multi-Agent Classic Board/Card Games & MPE | Full AEC & Parallel API Support (`tictactoe_v3`, `leduc_holdem_v4`) |
| **Atari (ALE / Shimmy)** | Pixel-based Visual Control | Supported via Atari Wrappers (`PongNoFrameskip-v4`) |
| **MicroRTS** (`gym_microrts`) | Multi-Discrete Real-Time Strategy | Supported via Multi-Discrete PPO |
| **MiniGrid** (`minigrid`) | Grid-world Navigation & POMDPs | Fully Compatible |

### 🧪 Demonstrated Environments in `examples/`

| Category | Environment | Example Scripts |
|---|---|---|
| **Discrete Control** | `CartPole-v1` | `examples/dqn/`, `examples/ppo/`, `examples/actor_critic/`, `examples/stream_rl/` |
| **Continuous Control** | `Pendulum-v1`, `HalfCheetah-v4` | `examples/ppo/ppo_pendulum.py`, `examples/ppo/ppo_mujoco.py`, `examples/actor_critic/` |
| **Visual Control** | `PongNoFrameskip-v4` | `examples/ppo/ppo_atari.py`, `examples/ppo/ppo_lstm_atari.py` |
| **Multi-Discrete** | `MicroRTS` | `examples/ppo/ppo_multidiscrete_microrts.py` |
| **Multi-Agent Board Games** | `tictactoe_v3` (PettingZoo) | `examples/alphazero/alphazero_tictactoe.py`, `alphazero_mp_tictactoe.py` |
| **Plasticity & Continual Learning** | Permuted MNIST, Bit-Flipping, ETTm2 | `examples/plasticity/`, `examples/stream_rl/` |
| **Meta-Optimization & Tracking** | Drifting Concept, Non-Stationary Random Walk | `examples/meta_optimization/idbd_feature_selection.py`, `auto_step.py` |
| **Tabular & Linear MDPs** | Sutton's Random Walk, Baird's Counterexample | `examples/td_learning/` |

---


## 💡 The Philosophy: Resolving the RL Paradigm War

Developing Reinforcement Learning algorithms typically forces researchers to choose between three flawed paradigms:

1. **The Monolith (Single-File Implementations):** 
   * *Pros:* Everything is in one local scope; logging, file-diffing, and fast prototyping are frictionless.
   * *Cons:* Combinatorial explosion of parameters and massive `if/else` configuration trees (e.g., handling frame stacking, continuous actions, recurrent states). Fixing a bug in one DQN implementation does not propagate to others, leading to destructive **feature drift**. Moving to multi-GPU (DDP) or TPU injections pollutes pure mathematical logic with system-level engineering blocks. Unit testing a loop monolith is practically impossible, forcing reliance on slow, flaky integration tests.
2. **Standard OOP & Strategy Patterns:**
   * *Pros:* Promotes abstraction and modular reuse.
   * *Cons:* Hides design details and heavily relies on internal state mutations (`self._step_count`, `self._hidden_state`), creating subtle off-by-one errors and tracking nightmares. To customize a single feature, you must master deep parent interfaces, inheriting a mountain of hidden knowledge debt. Mixing orthogonal features causes a combinatorial explosion of classes (e.g., `RecurrentContinuousPPO`), or God Classes. Strict interfaces strip away critical paper-specific mathematical optimizations to stay general.
3. **Execution Graphs & DAGs:**
   * *Pros:* Maximizes component reuse and layout validation.
   * *Cons:* Deep data-flow graphs aggressively reject standard Python dynamic control flows (like nested dynamic `while` loops inside an MCTS tree search), forcing clunky graph operators. Building a graph system creates immense engineering overhead, resulting in 10–15 node classes for a baseline algorithm, which damages execution speed in PyTorch due to structural dictionary-passing overhead.

### 🛠️ The Solution: Functional Core, Imperative Shell

This repository maps out a clean compromise. We isolate mathematical and algorithmic actions into a **Functional Core** composed of stateless, pure, side-effect-free functions. We then assemble these primitives inside an easy-to-read, linear, monolithic loop—the **Imperative Shell**. This allows our implementations to benefit from the monolithic paradigm, while allowing code reuse, modularity, and testing. Since each function is pure and has a simple interface, and (for the most part) functions don't use other functions internally, it is trivial to read and understand the codebase without getting lost in abstractions or having to dive deep into the codebase. 

```python
# --- 1. Initialization (Defining the State) ---
params = init_network()
optimizer_state = init_optimizer()
buffer_state = init_buffer(capacity=10000)
env_state, obs = env.reset()
hidden_state = init_rnn_state()

# --- 2. The Monolithic Loop (The Imperative Shell) ---
for step in range(MAX_STEPS):
    # 1. Act (Pure function)
    action, next_hidden_state = select_action(params, obs, hidden_state)
    
    # 2. Step Env (Pure-ish function adapter)
    next_env_state, next_obs, reward, done = env.step(env_state, action)
    
    # 3. Add to Buffer (Pure state mutation)
    transition = (obs, action, reward, hidden_state)
    buffer_state = add_to_buffer(buffer_state, transition)
    
    # Update loop states
    obs, env_state, hidden_state = next_obs, next_env_state, next_hidden_state
    
    # --- 3. The Functional Update Core ---
    if step % UPDATE_FREQ == 0:
        batch, rng_key = sample_buffer(buffer_state, rng_key, BATCH_SIZE)
        
        # Standard PyTorch Optimization
        optimizer.zero_grad(set_to_none=True)
        loss.backward()
        optimizer.step()
        
        # Monolithic layout makes logging and tracking effortless
        wandb.log({"loss": loss, "step": step})
```

## ⚡ Core Engineering Practices (TODO: improve this section, add more of our actual rules in CONTRIBUTING.md)
1. Explicit Over Implicit
We avoid magic configuration objects, automated parameter routing, and hidden global variables. High-level orchestration functions are stateless and clear. You pass tensors explicitly, enabling seamless tracking and debugging.

2. Fail Fast & Guardrails
We utilize strong type hints, strict validation checks, and inline assertions at the boundaries of the functional layers. Shape mismatches, device conflicts, and invalid boundaries throw exceptions at their origin point—not deep inside a backend compiled gradient execution.

3. Documentation by Signature
Functions are structured so that a researcher can easily interpret the underlying math just by inspecting the name, inputs, and type annotations, removing the need to trace internal source operations across 15 separate tracking scripts.

4. PyTorch Native Performance & Conventions
torch.compile Friendly: Pure functions minimize internal state and dictionary unpacks, enabling the compiler to run graph optimizations across the mathematical core.

No CUDA Synchronizations: Absolutely zero internal calls to .item(), .tolist(), or .numpy() inside the computational blocks. This keeps the CPU and GPU timelines decoupled, eliminating latency bubbles.

Device & Dtype Agnostic: We never hardcode strings like device='cuda'. When initializing tensors, we use factory methods like .new_zeros() or torch.zeros_like() relative to the incoming tensor states to protect against distributed data-parallel breaks.

Explicit Dimensions: PyTorch broadcasting can mask bugs. We favor explicit .unsqueeze() calls, allowing code readers to visually audit array matching.
