diff --git a/.github/pull_request_template.md b/.github/pull_request_template.md
index cf1796a..6fc8424 100644
--- a/.github/pull_request_template.md
+++ b/.github/pull_request_template.md
@@ -8,8 +8,7 @@
### Related Issues
-- Fixes: #
-- Related: #
+- #
### Screenshots / Notes
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
index 67683a3..f7d7840 100644
--- a/.github/workflows/ci.yml
+++ b/.github/workflows/ci.yml
@@ -21,12 +21,16 @@ jobs:
- name: Install dependencies
run: uv sync
- # Ambiguous reporting
- # - name: Run linting
- # run: uv run ruff check .
-
+ # Ambiguous reporting
+ # - name: Run linting
+ # run: uv run ruff check .
+
+ # Too slow
+ # - name: Run type checking
+ # run: uv run mypy src/
+
- name: Run type checking
- run: uv run mypy src/
+ run: uv run ty check .
- name: Run tests
run: uv run pytest
diff --git a/.gitignore b/.gitignore
index 77eaf4d..9f589ed 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,206 +1,36 @@
-# CU MIND
-# Files
-scratch.ipynb
-test.json
-test.json.bak
-
-# Directory
-checkpoints/
-logs/
-wandb/
-
-
-# Byte-compiled / optimized / DLL files
+# General files
+*.ipynb
+*.json
+scratch.*
+
+# Directories
+artifacts/
+experiments/
+gymnasium/
+openSpiel/
+
+# Python cache and build
__pycache__/
-*.py[cod]
-*$py.class
-
-# C extensions
-*.so
-
-# Distribution / packaging
-.Python
-build/
-develop-eggs/
dist/
-downloads/
-eggs/
-.eggs/
-lib/
-lib64/
-parts/
-sdist/
-var/
-wheels/
-share/python-wheels/
*.egg-info/
-.installed.cfg
-*.egg
-MANIFEST
-
-# PyInstaller
-# Usually these files are written by a python script from a template
-# before PyInstaller builds the exe, so as to inject date/other infos into it.
-*.manifest
-*.spec
-
-# Installer logs
-pip-log.txt
-pip-delete-this-directory.txt
-
-# Unit test / coverage reports
-htmlcov/
-.tox/
-.nox/
-.coverage
-.coverage.*
-.cache
-nosetests.xml
-coverage.xml
-*.cover
-*.py,cover
-.hypothesis/
-.pytest_cache/
-cover/
-
-# Translations
-*.mo
-*.pot
-
-# Django stuff:
-*.log
-local_settings.py
-db.sqlite3
-db.sqlite3-journal
-
-# Flask stuff:
-instance/
-.webassets-cache
-
-# Scrapy stuff:
-.scrapy
-
-# Sphinx documentation
-docs/_build/
-
-# PyBuilder
-.pybuilder/
-target/
-
-# Jupyter Notebook
-.ipynb_checkpoints
+build/
-# IPython
+# Jupyter & IPython
+.ipynb_checkpoints/
profile_default/
ipython_config.py
-# pyenv
-# For a library or package, you might want to ignore these files since the code is
-# intended to run in multiple environments; otherwise, check them in:
-# .python-version
-
-# pipenv
-# According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
-# However, in case of collaboration, if having platform-specific dependencies or dependencies
-# having no cross-platform support, pipenv may install dependencies that don't work, or not
-# install all needed dependencies.
-#Pipfile.lock
-
-# UV
-# Similar to Pipfile.lock, it is generally recommended to include uv.lock in version control.
-# This is especially recommended for binary packages to ensure reproducibility, and is more
-# commonly ignored for libraries.
-#uv.lock
-
-# poetry
-# Similar to Pipfile.lock, it is generally recommended to include poetry.lock in version control.
-# This is especially recommended for binary packages to ensure reproducibility, and is more
-# commonly ignored for libraries.
-# https://python-poetry.org/docs/basic-usage/#commit-your-poetrylock-file-to-version-control
-#poetry.lock
-
-# pdm
-# Similar to Pipfile.lock, it is generally recommended to include pdm.lock in version control.
-#pdm.lock
-# pdm stores project-wide configurations in .pdm.toml, but it is recommended to not include it
-# in version control.
-# https://pdm.fming.dev/latest/usage/project/#working-with-version-control
-.pdm.toml
-.pdm-python
-.pdm-build/
-
-# PEP 582; used by e.g. github.com/David-OConnor/pyflow and github.com/pdm-project/pdm
-__pypackages__/
-
-# Celery stuff
-celerybeat-schedule
-celerybeat.pid
-
-# SageMath parsed files
-*.sage.py
-
-# Environments
-.env
-.venv
+# Virtual environments
+.venv*/
env/
venv/
-ENV/
-env.bak/
-venv.bak/
-
-# Spyder project settings
-.spyderproject
-.spyproject
-
-# Rope project settings
-.ropeproject
-
-# mkdocs documentation
-/site
-# mypy
+# Quality Control
.mypy_cache/
-.dmypy.json
-dmypy.json
-
-# Pyre type checker
-.pyre/
-
-# pytype static type analyzer
-.pytype/
-
-# Cython debug symbols
-cython_debug/
-
-# PyCharm
-# JetBrains specific template is maintained in a separate JetBrains.gitignore that can
-# be found at https://github.com/github/gitignore/blob/main/Global/JetBrains.gitignore
-# and can be added to the global gitignore or merged into this file. For a more nuclear
-# option (not recommended) you can uncomment the following to ignore the entire idea folder.
-#.idea/
-
-# Abstra
-# Abstra is an AI-powered process automation framework.
-# Ignore directories containing user credentials, local state, and settings.
-# Learn more at https://abstra.io/docs
-.abstra/
-
-# Visual Studio Code
-# Visual Studio Code specific template is maintained in a separate VisualStudioCode.gitignore
-# that can be found at https://github.com/github/gitignore/blob/main/Global/VisualStudioCode.gitignore
-# and can be added to the global gitignore or merged into this file. However, if you prefer,
-# you could uncomment the following to ignore the enitre vscode folder
-# .vscode/
-
-# Ruff stuff:
.ruff_cache/
+.pytest_cache/
-# PyPI configuration file
-.pypirc
-
-# Cursor
-# Cursor is an AI-powered code editor. `.cursorignore` specifies files/directories to
-# exclude from AI features like autocomplete and code analysis. Recommended for sensitive data
-# refer to https://docs.cursor.com/context/ignore-files
-.cursorignore
-.cursorindexingignore
+# OS/Editor files
+.DS_Store
+*.swp
+*.bak
diff --git a/.vscode/launch.json b/.vscode/launch.json
index b720e6e..6ab65b5 100644
--- a/.vscode/launch.json
+++ b/.vscode/launch.json
@@ -3,10 +3,10 @@
"version": "0.2.0",
"configurations": [
{
- "name": "cartpole",
+ "name": "debug",
"type": "debugpy",
"request": "launch",
- "program": "${workspaceFolder}/src/cartpole.py",
+ "program": "${workspaceFolder}/src/main.py",
"console": "integratedTerminal",
"stopOnEntry": true
},
diff --git a/CuMind.svg b/CuMind.svg
new file mode 100644
index 0000000..f8f08b1
--- /dev/null
+++ b/CuMind.svg
@@ -0,0 +1,941 @@
+
+
+
+
+
diff --git a/Release_Notes.txt b/Release_Notes.txt
index 25221bd..d07d0f3 100644
--- a/Release_Notes.txt
+++ b/Release_Notes.txt
@@ -3,9 +3,17 @@ CuMind Release Notes
Document all technical changes introduced in this release as concise bullet points below.
+v0.2.0 (2025-08-09)
+----------------------
+Tarek Ibrahim ()
+- Unified experiment artifact management for consistent outputs.
+- Enhanced JAX support with typing and tracing capabilities.
+- Integrated type checking into CI and improved deployment workflows.
+- Includes critical bug fixes and performance improvements.
+
v0.1.95 (2025-08-03)
----------------------
-Tarek Ibrahim (68++)
+Tarek Ibrahim ()
- Added Weights & Biases (wandb) integration for experiment tracking
- Trainer improvements: debug mode, configurable number of batches, multi-backend support
- Added TqdmSink for progress bar logging
diff --git a/configuration.json b/configuration.json
index 34bbaa4..93a0146 100644
--- a/configuration.json
+++ b/configuration.json
@@ -40,7 +40,6 @@
"checkpoint_interval": 50,
"num_episodes": 2000,
"train_frequency": 2,
- "checkpoint_dir": "checkpoints",
"debug": false
},
"mcts": {
@@ -69,15 +68,25 @@
"wandb": false,
"title": "spamEggs",
"tags": ["foo", "bar"],
- "dir": "logs",
"level": "INFO",
"console": true,
"timestamps": true,
"tqdm": false
},
+ "tracing": {
+ "enabled": true,
+ "host_level": 1,
+ "python_level": 0,
+ "tpu_trace_mode": "TRACE_ONLY_XLA",
+ "tpu_num_sparse_cores_to_trace": 8,
+ "tpu_num_sparse_core_tiles_to_trace": 2,
+ "tpu_num_chips_to_profile_per_task": 8,
+ "save_memory": false
+ },
"device": "cpu",
"seed": 42,
- "validate": true,
- "multi_device": false
+ "multi_device": false,
+ "workspace": "artifacts",
+ "validate": true
}
}
diff --git a/cumind.svg b/cumind.svg
deleted file mode 100644
index a8ce2ee..0000000
--- a/cumind.svg
+++ /dev/null
@@ -1,529 +0,0 @@
-
-
-
-
-
diff --git a/gymnasium/box2d/.gitkeep b/gymnasium/box2d/.gitkeep
deleted file mode 100644
index e69de29..0000000
diff --git a/gymnasium/classic_control/Acrobot-v1/configuration.json b/gymnasium/classic_control/Acrobot-v1/configuration.json
deleted file mode 100644
index ae15860..0000000
--- a/gymnasium/classic_control/Acrobot-v1/configuration.json
+++ /dev/null
@@ -1,54 +0,0 @@
-{
- "Network architecture": {
- "hidden_dim": 128,
- "num_blocks": 2,
- "conv_channels": 32
- },
- "Training": {
- "batch_size": 64,
- "learning_rate": 0.01,
- "weight_decay": 0.0001,
- "target_update_frequency": 250,
- "checkpoint_interval": 50,
- "num_episodes": 1220,
- "train_frequency": 2,
- "checkpoint_root_dir": "checkpoints"
- },
- "MCTS": {
- "num_simulations": 25,
- "c_puct": 1.25,
- "dirichlet_alpha": 0.25,
- "exploration_fraction": 0.25
- },
- "Environment": {
- "env_name": "Acrobot-v1",
- "action_space_size": 3,
- "observation_shape": [
- 6
- ]
- },
- "Self-Play": {
- "num_unroll_steps": 5,
- "td_steps": 10,
- "discount": 0.997
- },
- "Memory": {
- "memory_capacity": 2000,
- "min_memory_size": 100,
- "min_memory_pct": 0.1,
- "per_alpha": 0.6,
- "per_epsilon": 1e-06,
- "per_beta": 0.4
- },
- "Data Types": {
- "model_dtype": "float32",
- "action_dtype": "int32",
- "target_dtype": "float32"
- },
- "Device": {
- "device_type": "cpu"
- },
- "Other": {
- "seed": 42
- }
-}
\ No newline at end of file
diff --git a/gymnasium/classic_control/CartPole-v1/configuration.json b/gymnasium/classic_control/CartPole-v1/configuration.json
deleted file mode 100644
index 79383e2..0000000
--- a/gymnasium/classic_control/CartPole-v1/configuration.json
+++ /dev/null
@@ -1,54 +0,0 @@
-{
- "Network architecture": {
- "hidden_dim": 128,
- "num_blocks": 2,
- "conv_channels": 32
- },
- "Training": {
- "batch_size": 64,
- "learning_rate": 0.01,
- "weight_decay": 0.0001,
- "target_update_frequency": 250,
- "checkpoint_interval": 50,
- "num_episodes": 1220,
- "train_frequency": 2,
- "checkpoint_root_dir": "checkpoints"
- },
- "MCTS": {
- "num_simulations": 25,
- "c_puct": 1.25,
- "dirichlet_alpha": 0.25,
- "exploration_fraction": 0.25
- },
- "Environment": {
- "env_name": "CartPole-v1",
- "action_space_size": 2,
- "observation_shape": [
- 4
- ]
- },
- "Self-Play": {
- "num_unroll_steps": 5,
- "td_steps": 10,
- "discount": 0.997
- },
- "Memory": {
- "memory_capacity": 2000,
- "min_memory_size": 100,
- "min_memory_pct": 0.1,
- "per_alpha": 0.6,
- "per_epsilon": 1e-06,
- "per_beta": 0.4
- },
- "Data Types": {
- "model_dtype": "float32",
- "action_dtype": "int32",
- "target_dtype": "float32"
- },
- "Device": {
- "device_type": "cpu"
- },
- "Other": {
- "seed": 42
- }
-}
\ No newline at end of file
diff --git a/gymnasium/classic_control/MountainCar-v0/configuration.json b/gymnasium/classic_control/MountainCar-v0/configuration.json
deleted file mode 100644
index cd94895..0000000
--- a/gymnasium/classic_control/MountainCar-v0/configuration.json
+++ /dev/null
@@ -1,54 +0,0 @@
-{
- "Network architecture": {
- "hidden_dim": 128,
- "num_blocks": 2,
- "conv_channels": 32
- },
- "Training": {
- "batch_size": 64,
- "learning_rate": 0.01,
- "weight_decay": 0.0001,
- "target_update_frequency": 250,
- "checkpoint_interval": 50,
- "num_episodes": 1220,
- "train_frequency": 2,
- "checkpoint_root_dir": "checkpoints"
- },
- "MCTS": {
- "num_simulations": 25,
- "c_puct": 1.25,
- "dirichlet_alpha": 0.25,
- "exploration_fraction": 0.25
- },
- "Environment": {
- "env_name": "MountainCar-v0",
- "action_space_size": 3,
- "observation_shape": [
- 2
- ]
- },
- "Self-Play": {
- "num_unroll_steps": 5,
- "td_steps": 10,
- "discount": 0.997
- },
- "Memory": {
- "memory_capacity": 2000,
- "min_memory_size": 100,
- "min_memory_pct": 0.1,
- "per_alpha": 0.6,
- "per_epsilon": 1e-06,
- "per_beta": 0.4
- },
- "Data Types": {
- "model_dtype": "float32",
- "action_dtype": "int32",
- "target_dtype": "float32"
- },
- "Device": {
- "device_type": "cpu"
- },
- "Other": {
- "seed": 42
- }
-}
\ No newline at end of file
diff --git a/gymnasium/classic_control/MountainCarContinuous-v0/configuration.json b/gymnasium/classic_control/MountainCarContinuous-v0/configuration.json
deleted file mode 100644
index 8955252..0000000
--- a/gymnasium/classic_control/MountainCarContinuous-v0/configuration.json
+++ /dev/null
@@ -1,56 +0,0 @@
-{
- "Network architecture": {
- "hidden_dim": 128,
- "num_blocks": 2,
- "conv_channels": 32
- },
- "Training": {
- "batch_size": 64,
- "learning_rate": 0.01,
- "weight_decay": 0.0001,
- "target_update_frequency": 250,
- "checkpoint_interval": 50,
- "num_episodes": 1220,
- "train_frequency": 2,
- "checkpoint_root_dir": "checkpoints"
- },
- "MCTS": {
- "num_simulations": 25,
- "c_puct": 1.25,
- "dirichlet_alpha": 0.25,
- "exploration_fraction": 0.25
- },
- "Environment": {
- "env_name": "MountainCarContinuous-v0",
- "action_space_size": [
- 1
- ],
- "observation_shape": [
- 2
- ]
- },
- "Self-Play": {
- "num_unroll_steps": 5,
- "td_steps": 10,
- "discount": 0.997
- },
- "Memory": {
- "memory_capacity": 2000,
- "min_memory_size": 100,
- "min_memory_pct": 0.1,
- "per_alpha": 0.6,
- "per_epsilon": 1e-06,
- "per_beta": 0.4
- },
- "Data Types": {
- "model_dtype": "float32",
- "action_dtype": "int32",
- "target_dtype": "float32"
- },
- "Device": {
- "device_type": "cpu"
- },
- "Other": {
- "seed": 42
- }
-}
\ No newline at end of file
diff --git a/gymnasium/classic_control/Pendulum-v1/configuration.json b/gymnasium/classic_control/Pendulum-v1/configuration.json
deleted file mode 100644
index 104bee8..0000000
--- a/gymnasium/classic_control/Pendulum-v1/configuration.json
+++ /dev/null
@@ -1,56 +0,0 @@
-{
- "Network architecture": {
- "hidden_dim": 128,
- "num_blocks": 2,
- "conv_channels": 32
- },
- "Training": {
- "batch_size": 64,
- "learning_rate": 0.01,
- "weight_decay": 0.0001,
- "target_update_frequency": 250,
- "checkpoint_interval": 50,
- "num_episodes": 1220,
- "train_frequency": 2,
- "checkpoint_root_dir": "checkpoints"
- },
- "MCTS": {
- "num_simulations": 25,
- "c_puct": 1.25,
- "dirichlet_alpha": 0.25,
- "exploration_fraction": 0.25
- },
- "Environment": {
- "env_name": "Pendulum-v1",
- "action_space_size": [
- 1
- ],
- "observation_shape": [
- 3
- ]
- },
- "Self-Play": {
- "num_unroll_steps": 5,
- "td_steps": 10,
- "discount": 0.997
- },
- "Memory": {
- "memory_capacity": 2000,
- "min_memory_size": 100,
- "min_memory_pct": 0.1,
- "per_alpha": 0.6,
- "per_epsilon": 1e-06,
- "per_beta": 0.4
- },
- "Data Types": {
- "model_dtype": "float32",
- "action_dtype": "int32",
- "target_dtype": "float32"
- },
- "Device": {
- "device_type": "cpu"
- },
- "Other": {
- "seed": 42
- }
-}
\ No newline at end of file
diff --git a/plots/episode_loss.png b/plots/episode_loss.png
deleted file mode 100644
index dee2107..0000000
Binary files a/plots/episode_loss.png and /dev/null differ
diff --git a/plots/policy_loss.png b/plots/policy_loss.png
deleted file mode 100644
index 7601168..0000000
Binary files a/plots/policy_loss.png and /dev/null differ
diff --git a/plots/reward.png b/plots/reward.png
deleted file mode 100644
index d634f1b..0000000
Binary files a/plots/reward.png and /dev/null differ
diff --git a/plots/reward_loss.png b/plots/reward_loss.png
deleted file mode 100644
index fad0de1..0000000
Binary files a/plots/reward_loss.png and /dev/null differ
diff --git a/plots/total_loss.png b/plots/total_loss.png
deleted file mode 100644
index cb2b791..0000000
Binary files a/plots/total_loss.png and /dev/null differ
diff --git a/plots/value_loss.png b/plots/value_loss.png
deleted file mode 100644
index d51a08e..0000000
Binary files a/plots/value_loss.png and /dev/null differ
diff --git a/pyproject.toml b/pyproject.toml
index cdb672d..445225f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -1,69 +1,61 @@
-# =========================
-# Project Metadata
-# =========================
[project]
-name = "cumind"
-version = "0.1.95"
+name = "CuMind"
+version = "0.2.0"
description = "A JAX-based CuMind is a JAX-based RL framework inspired by Google DeepMind. Achieve's superhuman performance in complex domains without pretraining nor prior knowledge of their rules."
readme = "README.md"
requires-python = ">=3.12"
-dependencies = [
+
+[build-system]
+requires = ["hatchling"]
+build-backend = "hatchling.build"
+
+[dependency-groups]
+main = [
"numpy>=2.3.0",
"gymnasium>=1.1.1",
"pygame>=2.6.1",
"tqdm>=4.67.1",
- "jax>=0.6.2",
+ "jax>=0.6.2", # Uses CPU by default; see config.py:385 for device selection.
"flax>=0.10.6",
"optax>=0.2.5",
"chex>=0.1.89",
"wandb>=0.20.1",
+]
+dev = [
+ "ruff>=0.11.13",
+ "pytest>=8.4.0",
+ "mypy>=1.16.0",
"ipython>=9.3.0",
"ipykernel>=6.29.5",
+ "pydeps>=3.0.1",
+ "ty>=0.0.1a17",
]
-# JAX Dependencies:
-# - jax: Core JAX library for numerical computing
-# - flax: Neural network library built on JAX
-# - optax: Gradient-based optimization library
-# - chex: Testing and assertion utilities for JAX
-
-# Note: JAX CPU version is used by default
-# LOOK AT CONFIG.PY LINE 385 VALIDATING DEVICE.
-
-
-# =========================
-# Build System
-# =========================
-[build-system]
-requires = ["hatchling"]
-build-backend = "hatchling.build"
-
-# =========================
-# Dependency Groups
-# =========================
-[dependency-groups]
-dev = ["ruff>=0.11.13", "pytest>=8.4.0", "mypy>=1.16.0"]
-
-# =========================
-# Tooling Configuration
-# =========================
-
-# Ruff configuration
[tool.ruff]
line-length = 320
target-version = "py312"
[tool.ruff.lint]
-select = ["E", "F", "I", "N", "W"] # Errors, pyFlakes, Import, Naming, Warnings
-ignore = ["E203", "E501"]
+select = [
+ "E", # pycodestyle errors
+ "F", # pyflakes errors
+ "I", # isort checks
+ "N", # pep8-naming
+ "W", # warnings
+]
+ignore = ["E501"] # Ignore line length errors
-# Mypy configuration
[tool.mypy]
python_version = "3.12"
strict = true
warn_return_any = true
warn_unused_configs = true
-# =========================
-# UV Configuration
-# =========================
+
+[tool.ty.src]
+include = ["src/"]
+exclude = []
+
+
+[tool.uv]
+default-groups = ["main", "dev"]
diff --git a/scripts/exportconfig.py b/scripts/exportconfig.py
index 6f81f87..03871eb 100755
--- a/scripts/exportconfig.py
+++ b/scripts/exportconfig.py
@@ -1,5 +1,5 @@
#!./.venv/bin/python3.12
-from cumind.utils.config import cfg
-
if __name__ == "__main__":
- cfg.save("configuration.json")
+ from cumind.utils.config import Configuration
+
+ Configuration()._to_json("configuration.json")
diff --git a/scripts/plot.py b/scripts/plot.py
deleted file mode 100755
index d99b557..0000000
--- a/scripts/plot.py
+++ /dev/null
@@ -1,128 +0,0 @@
-#!./.venv/bin/python3.12
-import datetime
-import os
-import re
-
-import matplotlib.pyplot as plt
-
-
-def parse_log_file(filepath):
- episode_re = re.compile(r"(?P