From bfde9473ebd394b79b25c8e967115b5373b80d51 Mon Sep 17 00:00:00 2001
From: Evan Palmer <evanp922@gmail.com>
Date: Tue, 1 Sep 2026 13:50:55 -0700
Subject: [PATCH 1/4] Add CI workflows and GitHub release pipeline

Adds three GitHub Actions workflows plus a pre-commit config:

- format.yml: pre-commit gate (house convention), Python 3.14
- ci.yml: Ruff lint/format check pinned to 0.16.5, and uv.lock freshness
- release.yml: on a `v*` tag push, verifies the tag matches the
  pyproject version, builds sdist+wheel, and creates a GitHub Release
  with auto-generated notes (tags containing `-` marked prerelease)
- dependabot.yml: weekly github-actions updates

Supporting changes:

- Move `jax[cuda]` into a `cuda` extra so the base install is CPU-only;
  document `lrax[cuda]` in the README
- Declare `tqdm`, which trainer.py imports but only resolved transitively
- Fix UP045 in envs/mjx.py, a docstring typo in trainer.py, and a
  codespell hit in the README; ignore `statics` (Equinox jargon)
- Extend .gitignore with dist/ and .ruff_cache/

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018a38TsYHunZKi4F8smG6JQ
---
 .github/dependabot.yml        |  6 +++++
 .github/workflows/ci.yml      | 40 +++++++++++++++++++++++++++++
 .github/workflows/format.yml  | 31 +++++++++++++++++++++++
 .github/workflows/release.yml | 47 +++++++++++++++++++++++++++++++++++
 .gitignore                    |  2 ++
 .pre-commit-config.yaml       | 27 ++++++++++++++++++++
 README.md                     | 12 ++++++++-
 lrax/common/envs/mjx.py       |  4 +--
 lrax/common/trainer.py        |  2 +-
 pyproject.toml                |  5 ++--
 uv.lock                       | 12 ++++++---
 11 files changed, 177 insertions(+), 11 deletions(-)
 create mode 100644 .github/dependabot.yml
 create mode 100644 .github/workflows/ci.yml
 create mode 100644 .github/workflows/format.yml
 create mode 100644 .github/workflows/release.yml
 create mode 100644 .pre-commit-config.yaml

diff --git a/.github/dependabot.yml b/.github/dependabot.yml
new file mode 100644
index 0000000..79fc83a
--- /dev/null
+++ b/.github/dependabot.yml
@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: github-actions
+    directory: "/"
+    schedule:
+      interval: "weekly"
diff --git a/.github/workflows/ci.yml b/.github/workflows/ci.yml
new file mode 100644
index 0000000..1035b83
--- /dev/null
+++ b/.github/workflows/ci.yml
@@ -0,0 +1,40 @@
+name: Continuous Integration
+
+on:
+  push:
+    branches:
+      - main
+  pull_request:
+  workflow_dispatch:
+
+jobs:
+  lint:
+    name: Lint
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v7
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+
+      - name: Run Ruff linter
+        run: uvx ruff@0.16.5 check --output-format=github
+
+      - name: Run Ruff formatter
+        run: uvx ruff@0.16.5 format --check
+
+  lock:
+    name: Lockfile
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v7
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+        with:
+          python-version: "3.14"
+
+      - name: Check that uv.lock is up to date
+        run: uv lock --check
diff --git a/.github/workflows/format.yml b/.github/workflows/format.yml
new file mode 100644
index 0000000..3d85e85
--- /dev/null
+++ b/.github/workflows/format.yml
@@ -0,0 +1,31 @@
+name: Formatting (pre-commit)
+
+on:
+  pull_request:
+  push:
+    branches:
+      - main
+  workflow_dispatch:
+
+jobs:
+  pre-commit:
+    name: Format
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v7
+
+      - name: Setup Python
+        uses: actions/setup-python@v7
+        with:
+          python-version: "3.14"
+
+      - name: Run pre-commit
+        uses: pre-commit/action@v3.0.1
+        id: precommit
+
+      - name: Upload pre-commit changes
+        if: failure() && steps.precommit.outcome == 'failure'
+        uses: rhaschke/upload-git-patch-action@main
+        with:
+          name: pre-commit
diff --git a/.github/workflows/release.yml b/.github/workflows/release.yml
new file mode 100644
index 0000000..79af7ca
--- /dev/null
+++ b/.github/workflows/release.yml
@@ -0,0 +1,47 @@
+name: Release
+
+on:
+  push:
+    tags:
+      - "v*"
+
+permissions:
+  contents: write
+
+jobs:
+  release:
+    name: Build and publish release
+    runs-on: ubuntu-latest
+    steps:
+      - name: Checkout repository
+        uses: actions/checkout@v7
+        with:
+          fetch-depth: 0
+
+      - name: Setup Python
+        uses: actions/setup-python@v7
+        with:
+          python-version: "3.14"
+
+      - name: Install uv
+        uses: astral-sh/setup-uv@v7
+
+      - name: Verify tag matches project version
+        run: |
+          version=$(python -c 'import tomllib, pathlib; print(tomllib.loads(pathlib.Path("pyproject.toml").read_text())["project"]["version"])')
+          if [ "v${version}" != "${GITHUB_REF_NAME}" ]; then
+            echo "::error::Tag ${GITHUB_REF_NAME} does not match pyproject.toml version v${version}"
+            exit 1
+          fi
+
+      - name: Build distributions
+        run: uv build
+
+      - name: Create GitHub release
+        env:
+          GH_TOKEN: ${{ github.token }}
+        run: |
+          prerelease=""
+          case "${GITHUB_REF_NAME}" in *-*) prerelease="--prerelease" ;; esac
+          gh release create "${GITHUB_REF_NAME}" dist/* \
+            --generate-notes --verify-tag ${prerelease}
diff --git a/.gitignore b/.gitignore
index 670a936..a732df2 100644
--- a/.gitignore
+++ b/.gitignore
@@ -1,2 +1,4 @@
 __pycache__/
 .venv/
+.ruff_cache/
+dist/
diff --git a/.pre-commit-config.yaml b/.pre-commit-config.yaml
new file mode 100644
index 0000000..6497a5c
--- /dev/null
+++ b/.pre-commit-config.yaml
@@ -0,0 +1,27 @@
+repos:
+  - repo: https://github.com/astral-sh/ruff-pre-commit
+    rev: v0.16.5
+    hooks:
+      - id: ruff-check
+        args: ["--fix", "--exit-non-zero-on-fix"]
+      - id: ruff-format
+
+  - repo: https://github.com/codespell-project/codespell
+    rev: v2.4.1
+    hooks:
+      - id: codespell
+        args: ["-L", "statics"]
+
+  - repo: https://github.com/pre-commit/pre-commit-hooks
+    rev: v6.0.0
+    hooks:
+      - id: check-added-large-files
+      - id: check-case-conflict
+      - id: check-toml
+      - id: check-yaml
+      - id: check-merge-conflict
+      - id: debug-statements
+      - id: detect-private-key
+      - id: end-of-file-fixer
+      - id: mixed-line-ending
+      - id: trailing-whitespace
diff --git a/README.md b/README.md
index 440bce9..d8c8b3c 100644
--- a/README.md
+++ b/README.md
@@ -25,12 +25,21 @@ uv pip install git+https://github.com/OSU-LRAM/lrax.git
 uv add git+ssh://git@github.com/OSU-LRAM/lrax.git
 ```
 
+The base install pulls in CPU-only JAX. For CUDA-accelerated JAX on Linux,
+include the optional `cuda` dependencies
+
+```bash
+uv add "lrax[cuda] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
+```
+
 To install the MuJoCo helpers, include the optional `mjx` dependencies
 
 ```bash
 uv add "lrax[mjx] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
 ```
 
+Extras combine, e.g. `lrax[cuda,mjx]`.
+
 ## Usage
 
 See the following example for training a policy with PPO against a custom
@@ -80,9 +89,10 @@ class PointMassEnv(AbstractEnv):
             done=done,
             terminated=jnp.zeros(self.num_envs, dtype=bool),
             terminal_obs=obs,
-            aux={}, # include auxilliary data from the environment
+            aux={},  # include auxiliary data from the environment
         )
 
+
 key = jr.key(0)
 actor_key, critic_key, train_key = jr.split(key, 3)
 env = PointMassEnv()
diff --git a/lrax/common/envs/mjx.py b/lrax/common/envs/mjx.py
index 78cfba8..53ff6ba 100644
--- a/lrax/common/envs/mjx.py
+++ b/lrax/common/envs/mjx.py
@@ -18,8 +18,6 @@
 # OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
 # THE SOFTWARE.
 
-from typing import Optional
-
 import mujoco
 from jaxtyping import Array
 from mujoco import mjx
@@ -41,7 +39,7 @@ def load_model(path: str) -> Model:
     return mjx.put_model(model)
 
 
-def reset(model: Model, qpos: Optional[Array] = None) -> Data:
+def reset(model: Model, qpos: Array | None = None) -> Data:
     """Create an initial simulation state.
 
     Parameters
diff --git a/lrax/common/trainer.py b/lrax/common/trainer.py
index d4e2825..8132d65 100644
--- a/lrax/common/trainer.py
+++ b/lrax/common/trainer.py
@@ -266,7 +266,7 @@ def learn(
             `algorithm.init` once, then `algorithm.step` in a loop, threading the
             algorithm state `init` returns through every `step` call.
         - `optim`: The optax optimizer(s) used to update `model`. For algorithms that
-            require multiple optimzers (like `SHAC`), this will be a PyTree of
+            require multiple optimizers (like `SHAC`), this will be a PyTree of
             optimizers.
         - `num_iterations`: The number of times to call `algorithm.step`. (Keyword
             only argument.)
diff --git a/pyproject.toml b/pyproject.toml
index efc14c0..99ce52f 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -10,18 +10,19 @@ readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
     "equinox>=0.13.8",
-    "jax[cuda]>=0.10.2; sys_platform != 'darwin'",
-    "jax>=0.10.2; sys_platform == 'darwin'",
+    "jax>=0.10.2",
     "jaxtyping>=0.3.11",
     "lightning>=2.6.5",
     "matplotlib>=3.11.1",
     "optax>=0.2.8",
     "optimistix>=0.1.0",
     "pyyaml>=6.0.3",
+    "tqdm>=4.67.0",
     "wandb>=0.19.0",
 ]
 
 [project.optional-dependencies]
+cuda = ["jax[cuda]>=0.10.2; sys_platform != 'darwin'"]
 mjx = ["mujoco-mjx>=3.10.0"]
 
 [tool.ruff.lint]
diff --git a/uv.lock b/uv.lock
index a3d15fa..dd8f626 100644
--- a/uv.lock
+++ b/uv.lock
@@ -864,17 +864,20 @@ source = { editable = "." }
 dependencies = [
     { name = "equinox" },
     { name = "jax" },
-    { name = "jax", extra = ["cuda"], marker = "sys_platform != 'darwin'" },
     { name = "jaxtyping" },
     { name = "lightning" },
     { name = "matplotlib" },
     { name = "optax" },
     { name = "optimistix" },
     { name = "pyyaml" },
+    { name = "tqdm" },
     { name = "wandb" },
 ]
 
 [package.optional-dependencies]
+cuda = [
+    { name = "jax", extra = ["cuda"], marker = "sys_platform != 'darwin'" },
+]
 mjx = [
     { name = "mujoco-mjx" },
 ]
@@ -882,8 +885,8 @@ mjx = [
 [package.metadata]
 requires-dist = [
     { name = "equinox", specifier = ">=0.13.8" },
-    { name = "jax", marker = "sys_platform == 'darwin'", specifier = ">=0.10.2" },
-    { name = "jax", extras = ["cuda"], marker = "sys_platform != 'darwin'", specifier = ">=0.10.2" },
+    { name = "jax", specifier = ">=0.10.2" },
+    { name = "jax", extras = ["cuda"], marker = "sys_platform != 'darwin' and extra == 'cuda'", specifier = ">=0.10.2" },
     { name = "jaxtyping", specifier = ">=0.3.11" },
     { name = "lightning", specifier = ">=2.6.5" },
     { name = "matplotlib", specifier = ">=3.11.1" },
@@ -891,9 +894,10 @@ requires-dist = [
     { name = "optax", specifier = ">=0.2.8" },
     { name = "optimistix", specifier = ">=0.1.0" },
     { name = "pyyaml", specifier = ">=6.0.3" },
+    { name = "tqdm", specifier = ">=4.67.0" },
     { name = "wandb", specifier = ">=0.19.0" },
 ]
-provides-extras = ["mjx"]
+provides-extras = ["cuda", "mjx"]
 
 [[package]]
 name = "markupsafe"

From 6219ba482a39013663e6004965cbf25dc347d616 Mon Sep 17 00:00:00 2001
From: Evan Palmer <evanp922@gmail.com>
Date: Tue, 1 Sep 2026 13:55:19 -0700
Subject: [PATCH 2/4] Remove unused MJX helpers

lrax/common/envs/mjx.py was three thin pass-throughs over mujoco.mjx
(load_model / reset / step), not re-exported from envs/__init__.py, not
imported anywhere, and with signatures that don't match the AbstractEnv
interface. Drop the module, the `mjx` optional extra, and the README
section. Reintroduce a `mujoco-mjx` dependency when there is an actual
AbstractEnv subclass backed by MJX to justify it.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018a38TsYHunZKi4F8smG6JQ
---
 README.md               |   8 ---
 lrax/common/envs/mjx.py |  76 --------------------------
 pyproject.toml          |   1 -
 uv.lock                 | 117 +---------------------------------------
 4 files changed, 1 insertion(+), 201 deletions(-)
 delete mode 100644 lrax/common/envs/mjx.py

diff --git a/README.md b/README.md
index d8c8b3c..510657b 100644
--- a/README.md
+++ b/README.md
@@ -32,14 +32,6 @@ include the optional `cuda` dependencies
 uv add "lrax[cuda] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
 ```
 
-To install the MuJoCo helpers, include the optional `mjx` dependencies
-
-```bash
-uv add "lrax[mjx] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
-```
-
-Extras combine, e.g. `lrax[cuda,mjx]`.
-
 ## Usage
 
 See the following example for training a policy with PPO against a custom
diff --git a/lrax/common/envs/mjx.py b/lrax/common/envs/mjx.py
deleted file mode 100644
index 53ff6ba..0000000
--- a/lrax/common/envs/mjx.py
+++ /dev/null
@@ -1,76 +0,0 @@
-# Copyright 2026, Laboratory for Robotics and Applied Mechanics
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
-# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-
-import mujoco
-from jaxtyping import Array
-from mujoco import mjx
-from mujoco.mjx import Data, Model
-
-
-def load_model(path: str) -> Model:
-    """Load an `mjx.Model` from a configuration file.
-
-    Parameters
-    ----------
-    - `path`: The full path to the `xml` configuration file.
-
-    Returns
-    -------
-    The loaded `mjx.Model`.
-    """
-    model = mujoco.MjModel.from_xml_path(path)  # type: ignore
-    return mjx.put_model(model)
-
-
-def reset(model: Model, qpos: Array | None = None) -> Data:
-    """Create an initial simulation state.
-
-    Parameters
-    ----------
-    - `model`: The MJX model to create the state for.
-    - `qpos`: The initial configuration.
-
-    Returns
-    -------
-    The initial simulation state.
-    """
-    data = mjx.make_data(model)
-
-    if qpos is not None:
-        data = data.replace(qpos=qpos)
-
-    return mjx.forward(model, data)
-
-
-def step(model: Model, data: Data, control: Array) -> Data:
-    """Apply a control to the system and advance the simulation by one timestep.
-
-    Parameters
-    ----------
-    - `model`: The MJX model being simulated.
-    - `data`: The current simulation state.
-    - `control`: A JAX array representing the control input to apply to the system.
-
-    Returns
-    -------
-    The next simulation state.
-    """
-    data = data.replace(ctrl=control)
-    return mjx.step(model, data)
diff --git a/pyproject.toml b/pyproject.toml
index 99ce52f..1406d04 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -23,7 +23,6 @@ dependencies = [
 
 [project.optional-dependencies]
 cuda = ["jax[cuda]>=0.10.2; sys_platform != 'darwin'"]
-mjx = ["mujoco-mjx>=3.10.0"]
 
 [tool.ruff.lint]
 ignore = ["F722", "F821", "E731", "F811"]
diff --git a/uv.lock b/uv.lock
index dd8f626..81a110f 100644
--- a/uv.lock
+++ b/uv.lock
@@ -413,22 +413,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/ae/d6/69a76c8ccdef14af687c497040292a46e59fc7a0ab24724b60e50ca61030/equinox-0.13.8-py3-none-any.whl", hash = "sha256:ca004348533cc30a63ebe8823d7dd4bb626dce17743d40bbddb89b402ef2a240", size = 185813, upload-time = "2026-05-05T10:03:41.673Z" },
 ]
 
-[[package]]
-name = "etils"
-version = "1.14.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/26/ce/6e067242fde898841922ac6fc82b0bb2fe35c38e995880bdffdfbe30182a/etils-1.14.0.tar.gz", hash = "sha256:8136e7f4c4173cd0af0ca5481c4475152f0b8686192951eefa60ee8711e1ede4", size = 108127, upload-time = "2026-03-04T17:41:36.291Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5a/3d/589663aeeacd59bb2f3e8596bfd3e81cf0fb18d70bb433199041f469771b/etils-1.14.0-py3-none-any.whl", hash = "sha256:b5df7341f54dbe1405a4450b2741207b4a8c279780402b45f87202b94dfc52b4", size = 172934, upload-time = "2026-03-04T17:41:35.01Z" },
-]
-
-[package.optional-dependencies]
-epath = [
-    { name = "fsspec" },
-    { name = "typing-extensions" },
-    { name = "zipp" },
-]
-
 [[package]]
 name = "filelock"
 version = "3.32.2"
@@ -582,22 +566,6 @@ http = [
     { name = "aiohttp" },
 ]
 
-[[package]]
-name = "glfw"
-version = "2.10.2"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/0f/62/096058bcb4b4fb28f7ecd28fb048f07969d90b243c417af5f6d09d45a0c2/glfw-2.10.2.tar.gz", hash = "sha256:5d2cf97c66bc42a6b583be0e307eae5a3945438322e2ed0c5e4f14dc251d693d", size = 36307, upload-time = "2026-07-21T14:42:36.732Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/c4/80/5e575ec14f54dc0a644e7215985b9ee7c5044fcc3594d6b83dfdeb04ccd4/glfw-2.10.2-py2.py3-none-macosx_10_6_intel.whl", hash = "sha256:a04d25bb535a3a29e173ea1abd658e161681b5d3816e00bc51c122f74cd7fa3a", size = 110295, upload-time = "2026-07-21T14:42:25.625Z" },
-    { url = "https://files.pythonhosted.org/packages/e5/2e/22216db429690aa51fa87d2e2b5e6b610c5a37b9df7768da2927f52f8704/glfw-2.10.2-py2.py3-none-macosx_11_0_arm64.whl", hash = "sha256:c000b8a5e5fb4374b2c18d12347255ccb92001b9bcf80ef74c56072acbb98fe3", size = 107146, upload-time = "2026-07-21T14:42:26.854Z" },
-    { url = "https://files.pythonhosted.org/packages/75/a4/9ccee9d5c3b9c5d6bafce3e8d53ab8318eff30124e3c19b95580bebaa6d4/glfw-2.10.2-py2.py3-none-manylinux2014_aarch64.whl", hash = "sha256:0ffe2b48b51f37b05e8f027f57c0b8cd213affaf22fc83401bd9621692d4c6c8", size = 235005, upload-time = "2026-07-21T14:42:27.955Z" },
-    { url = "https://files.pythonhosted.org/packages/77/44/37aeed50c581c76e1c5bb83b696a06f98f8ecc979ed8564dde3eef908e8b/glfw-2.10.2-py2.py3-none-manylinux2014_x86_64.whl", hash = "sha256:82daf1f87b2a48815637dc6f6720d4a55a57ba1d3683322dc20dc28065754f97", size = 246951, upload-time = "2026-07-21T14:42:29.107Z" },
-    { url = "https://files.pythonhosted.org/packages/52/fc/29b4f1a8c8e33d8934781330a0845458b01da7f389d9fe6a3350e752b7f3/glfw-2.10.2-py2.py3-none-manylinux_2_28_aarch64.whl", hash = "sha256:7604425ec95665cc6a25c3ed7d07e9660d394a34472ceffaec0843c844e7b649", size = 236018, upload-time = "2026-07-21T14:42:30.623Z" },
-    { url = "https://files.pythonhosted.org/packages/69/84/1b98671314ea5f40397586b6cd14db913de3196333be558fdc177e61708f/glfw-2.10.2-py2.py3-none-manylinux_2_28_x86_64.whl", hash = "sha256:b860de3ca0686182483f98f3ddd12e660acf25b3e0d521450ec9a3f999f72a65", size = 248490, upload-time = "2026-07-21T14:42:32.377Z" },
-    { url = "https://files.pythonhosted.org/packages/12/d9/bd2e6c8dbadcf69fcd2289705e75881258b46221c196903a52e9f6e97322/glfw-2.10.2-py2.py3-none-win32.whl", hash = "sha256:a94aefe8c48886fd83cd6e11e936846166a4b5b94abf1f4bf599d984774b0b1b", size = 557654, upload-time = "2026-07-21T14:42:33.903Z" },
-    { url = "https://files.pythonhosted.org/packages/e0/6c/4ca5f3ab85a8d7f612ce857208726e9b834e54e559751e3d4f98cc2f7509/glfw-2.10.2-py2.py3-none-win_amd64.whl", hash = "sha256:15fd0666cd8f1b0ecb535c3abb712b8ddda053bf8d16de2353b1c6ced0e65402", size = 564433, upload-time = "2026-07-21T14:42:35.465Z" },
-]
-
 [[package]]
 name = "idna"
 version = "3.18"
@@ -878,9 +846,6 @@ dependencies = [
 cuda = [
     { name = "jax", extra = ["cuda"], marker = "sys_platform != 'darwin'" },
 ]
-mjx = [
-    { name = "mujoco-mjx" },
-]
 
 [package.metadata]
 requires-dist = [
@@ -890,14 +855,13 @@ requires-dist = [
     { name = "jaxtyping", specifier = ">=0.3.11" },
     { name = "lightning", specifier = ">=2.6.5" },
     { name = "matplotlib", specifier = ">=3.11.1" },
-    { name = "mujoco-mjx", marker = "extra == 'mjx'", specifier = ">=3.10.0" },
     { name = "optax", specifier = ">=0.2.8" },
     { name = "optimistix", specifier = ">=0.1.0" },
     { name = "pyyaml", specifier = ">=6.0.3" },
     { name = "tqdm", specifier = ">=4.67.0" },
     { name = "wandb", specifier = ">=0.19.0" },
 ]
-provides-extras = ["cuda", "mjx"]
+provides-extras = ["cuda"]
 
 [[package]]
 name = "markupsafe"
@@ -1061,55 +1025,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/43/e3/7d92a15f894aa0c9c4b49b8ee9ac9850d6e63b03c9c32c0367a13ae62209/mpmath-1.3.0-py3-none-any.whl", hash = "sha256:a0b2b9fe80bbcd81a6647ff13108738cfb482d481d826cc0e02f5b35e5c88d2c", size = 536198, upload-time = "2023-03-07T16:47:09.197Z" },
 ]
 
-[[package]]
-name = "mujoco"
-version = "3.11.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "absl-py" },
-    { name = "etils", extra = ["epath"] },
-    { name = "glfw" },
-    { name = "numpy" },
-    { name = "pyopengl" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/3b/2d/290f3f4062eec1c0d5246de0a75f1b18b59a1961081f49ddc97333fc8263/mujoco-3.11.0.tar.gz", hash = "sha256:390856102af9547dfd87cb2791eb105a3d2e00a37323fe9a12ed17e512aff1a6", size = 1139880, upload-time = "2026-07-28T01:23:32.252Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/87/b3/b2b449b5978bcb13277c9e50d764e6d8cd9534f58f071fb1647724123e2c/mujoco-3.11.0-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:3609c198de090c8218b9cc62ca6133f0feede8edd103b743b8002f3f735fcd9b", size = 17511145, upload-time = "2026-07-28T01:22:43.668Z" },
-    { url = "https://files.pythonhosted.org/packages/1a/65/d8130bb8a673f9cdc8a8fb2ef02f9b6931708567de57f17395106e83b4e3/mujoco-3.11.0-cp312-cp312-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:517bff03623c3329a563f575afd7d731b0be5c4f92a54dcec934b99120b4a9ec", size = 17704436, upload-time = "2026-07-28T01:22:47.05Z" },
-    { url = "https://files.pythonhosted.org/packages/41/d6/7b79f5d8fbd019658a3b5feb6ffd09c1e727eaf93f518685d3cc105a28f5/mujoco-3.11.0-cp312-cp312-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:f40214fefc8c2fe0002a3c8abadf30de7a3330634f3c45cc29b407b40a0173fc", size = 18868030, upload-time = "2026-07-28T01:22:50.494Z" },
-    { url = "https://files.pythonhosted.org/packages/19/46/f994cd7d973c4db4f6b5af19ba6eb62703b9aafc8f894c5bc5ceadd31b0d/mujoco-3.11.0-cp312-cp312-win_amd64.whl", hash = "sha256:98c08a5eec9411ecd6f763f1e635fbc8f31e3ea3701584d6b7edcc24d8d6c575", size = 15791637, upload-time = "2026-07-28T01:22:53.922Z" },
-    { url = "https://files.pythonhosted.org/packages/d7/16/4b822d608cb4bc5aafe910b6f6947c73ecd467470935408e4ad0d207023e/mujoco-3.11.0-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:dd282aae46dc7350472bd4a81dd789ea011ff961c00fd60c9434e67bfd0519a3", size = 17511474, upload-time = "2026-07-28T01:22:57.111Z" },
-    { url = "https://files.pythonhosted.org/packages/47/73/4cbe741986ed86f13e86ad92a11de870bdec4979eaafcbc7ef3164680260/mujoco-3.11.0-cp313-cp313-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:b1ea2ed1d20b7cca55fc1e0257f2d4f6ab6e83855f669530fb3f55c85da80aac", size = 17704381, upload-time = "2026-07-28T01:23:00.158Z" },
-    { url = "https://files.pythonhosted.org/packages/90/2e/843071e21831fa5ef4029de6874789f9fd32f340d42fced383757c484bdd/mujoco-3.11.0-cp313-cp313-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:3379c36b508474127b08e51e77942c493c2f1d7215b425e3b95de08dcba662e0", size = 18868301, upload-time = "2026-07-28T01:23:03.054Z" },
-    { url = "https://files.pythonhosted.org/packages/bb/25/2a518a33b359b0798a9e8b55f30a4c9ab0cadd0cab8d85872ff9da6ca309/mujoco-3.11.0-cp313-cp313-win_amd64.whl", hash = "sha256:d48b4a18d409ce4c746f32b3b402a8419b7450d9a89de68b81d73a3ef333d0e4", size = 15792231, upload-time = "2026-07-28T01:23:05.798Z" },
-    { url = "https://files.pythonhosted.org/packages/36/56/1d434ab8e072a6ca9074a27951cc08770db50daa4f2908fc51c8fe1c73dd/mujoco-3.11.0-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:d0f031b5b747a96c07938a47305f46079b75697282723e3ae1864c089bb56c02", size = 17525650, upload-time = "2026-07-28T01:23:08.562Z" },
-    { url = "https://files.pythonhosted.org/packages/5f/5a/f2fce0b815a51a3cb87d62e0fb3647f8d10330b30a8e7d381b733773263e/mujoco-3.11.0-cp314-cp314-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:eeb372621f885a0ad16ef879c91ff74d47b3b9533babdc16e07303836a23d27e", size = 17720735, upload-time = "2026-07-28T01:23:11.205Z" },
-    { url = "https://files.pythonhosted.org/packages/f9/98/fad81239c7b827a64a472206a6bff5346e3fa949e6ce4945595329d3fdc7/mujoco-3.11.0-cp314-cp314-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:06af29af1fa40156126165f2634481ef7c0f8dabb9f07892c5af30039c5a59dc", size = 18866078, upload-time = "2026-07-28T01:23:13.982Z" },
-    { url = "https://files.pythonhosted.org/packages/81/33/93ba5542249502ab9d7c1d1f3e38cd915a0259bc4c8cd124c05dbe17e7bb/mujoco-3.11.0-cp314-cp314-win_amd64.whl", hash = "sha256:2a9bf53284604397b7cd041cecb22c8f6b10afdddb4c723d5ff493f9f6b630d7", size = 16436726, upload-time = "2026-07-28T01:23:29.6Z" },
-    { url = "https://files.pythonhosted.org/packages/5b/b8/bc0835e43ed36ab465b4c3b1f6f8324c025207acc4d1e9c24ffaa4ebfb87/mujoco-3.11.0-cp314-cp314t-macosx_11_0_arm64.whl", hash = "sha256:b2be7535332976a82ec69a118cdeb3db3580aba19e3d82bca91c95a416e2a9ce", size = 17696886, upload-time = "2026-07-28T01:23:17.225Z" },
-    { url = "https://files.pythonhosted.org/packages/36/a4/8444786ee44588a02842069eeaca31fc1a25357ac9b5ef03d96e424dc8fd/mujoco-3.11.0-cp314-cp314t-manylinux_2_27_aarch64.manylinux_2_28_aarch64.whl", hash = "sha256:80cf42c0153d8dac61fb56724ba3fd8b0e0fb7093c013d57ae9e51f894262338", size = 17952805, upload-time = "2026-07-28T01:23:20.432Z" },
-    { url = "https://files.pythonhosted.org/packages/7b/bd/46a43763dac53e840a9ce1f204f9d80fb07fc7c2a6805cb023fdc35a0986/mujoco-3.11.0-cp314-cp314t-manylinux_2_27_x86_64.manylinux_2_28_x86_64.whl", hash = "sha256:8205329fce598bc3aede8ef530da6a70c529c8f4a91efdc79933614ba3f126fb", size = 19029401, upload-time = "2026-07-28T01:23:23.656Z" },
-    { url = "https://files.pythonhosted.org/packages/ad/d8/10200a2c1dc65738bc21375b23fe986e4bf435be045ea91eaf698e86cfd7/mujoco-3.11.0-cp314-cp314t-win_amd64.whl", hash = "sha256:0804a70ccd9b669ed36e1ede45781d0d707f2940d560ed756647cc8d9ade847c", size = 16460526, upload-time = "2026-07-28T01:23:26.81Z" },
-]
-
-[[package]]
-name = "mujoco-mjx"
-version = "3.11.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "absl-py" },
-    { name = "etils", extra = ["epath"] },
-    { name = "jax" },
-    { name = "jaxlib" },
-    { name = "mujoco" },
-    { name = "scipy" },
-    { name = "trimesh" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/c1/94/fb86b46d0d046d98867837ee4c73c7c69011397a7ccfe5eecd492c3355b8/mujoco_mjx-3.11.0.tar.gz", hash = "sha256:40395d6293d93aa0a2ae11fa3d0c59236754ec923e40093c5251db719cc425cc", size = 7067359, upload-time = "2026-07-28T01:26:27.834Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/5a/17/420535e6e19c749b608ce7ed46a89fee966b2c7ef62bbdce80470dab25f3/mujoco_mjx-3.11.0-py3-none-any.whl", hash = "sha256:5c3a96a927717a2fff89570da241afc7fe386d315c68a8e37d72a7f1549db846", size = 7163512, upload-time = "2026-07-28T01:26:25.743Z" },
-]
-
 [[package]]
 name = "multidict"
 version = "6.7.1"
@@ -1886,15 +1801,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/fa/c3/7c8b240552251faf6b3a957db200fcfbbcec36763c050428b601e0c9b83b/pydantic_core-2.46.4-graalpy312-graalpy250_312_native-manylinux_2_17_x86_64.manylinux2014_x86_64.whl", hash = "sha256:00c603d540afdd6b80eb39f078f33ebd46211f02f33e34a32d9f053bba711de0", size = 2147590, upload-time = "2026-05-06T13:39:29.883Z" },
 ]
 
-[[package]]
-name = "pyopengl"
-version = "3.1.10"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/6f/16/912b7225d56284859cd9a672827f18be43f8012f8b7b932bc4bd959a298e/pyopengl-3.1.10.tar.gz", hash = "sha256:c4a02d6866b54eb119c8e9b3fb04fa835a95ab802dd96607ab4cdb0012df8335", size = 1915580, upload-time = "2025-08-18T02:33:01.76Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/de/e4/1ba6f44e491c4eece978685230dde56b14d51a0365bc1b774ddaa94d14cd/pyopengl-3.1.10-py3-none-any.whl", hash = "sha256:794a943daced39300879e4e47bd94525280685f42dbb5a998d336cfff151d74f", size = 3194996, upload-time = "2025-08-18T02:32:59.902Z" },
-]
-
 [[package]]
 name = "pyparsing"
 version = "3.3.2"
@@ -2156,18 +2062,6 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/f9/1c/01bfd571a64e7f270e6bab5e33777debe0edc56759233ce84f27dec92d14/tqdm-4.70.0-py3-none-any.whl", hash = "sha256:7f585706bfddbdebf89daac705b2dfcc16890130727d3197ca62c732b4310953", size = 80184, upload-time = "2026-07-27T11:33:13.167Z" },
 ]
 
-[[package]]
-name = "trimesh"
-version = "5.0.0"
-source = { registry = "https://pypi.org/simple" }
-dependencies = [
-    { name = "numpy" },
-]
-sdist = { url = "https://files.pythonhosted.org/packages/a8/45/5b9943631f265073bb96ca4cc2a9de05d151cd8c1a6896fcbaac132a1355/trimesh-5.0.0.tar.gz", hash = "sha256:0195003198baaf2550aebe612254ba2f01c385cec46c37ecc5beca8291f79a13", size = 864019, upload-time = "2026-08-01T00:05:25.042Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/37/82/19a03ba344ecb66ea8caab697b3059e0fbea576420f99945944c479caa78/trimesh-5.0.0-py3-none-any.whl", hash = "sha256:51ec67d7f9f74b918f2a695da5fd511b9c084e96307648ae83f45b58f13f8009", size = 745494, upload-time = "2026-08-01T00:05:23.172Z" },
-]
-
 [[package]]
 name = "triton"
 version = "3.7.1"
@@ -2331,12 +2225,3 @@ wheels = [
     { url = "https://files.pythonhosted.org/packages/bf/f4/ed5c402ac8fde4403ed3366c2716bfddc8a6677ebd59f3d62772cc7fe468/yarl-1.24.5-cp314-cp314t-win_arm64.whl", hash = "sha256:cf139c02f5f23ef6532040a30ff662c00a318c952334f211046b8e60b7f17688", size = 97222, upload-time = "2026-07-20T02:07:41.55Z" },
     { url = "https://files.pythonhosted.org/packages/61/02/962c1cbfc401a30c1d034dc67ff395f64b52302c6d62de556c1fca99acc0/yarl-1.24.5-py3-none-any.whl", hash = "sha256:a33700d13d9b7d84fd10947b09ff69fb9a792e519c8cb9764a3ca70baa6c23a7", size = 58612, upload-time = "2026-07-20T02:07:43.461Z" },
 ]
-
-[[package]]
-name = "zipp"
-version = "4.1.0"
-source = { registry = "https://pypi.org/simple" }
-sdist = { url = "https://files.pythonhosted.org/packages/b9/d8/eab98a517c14134c0b2eb4e2387bc5f457334293ec5d2dd3857ec2966802/zipp-4.1.0.tar.gz", hash = "sha256:4cb57381f544315db7688e976e922a2b18cdb513d21cc194eb42232ba2a3e602", size = 26214, upload-time = "2026-05-18T20:08:57.967Z" }
-wheels = [
-    { url = "https://files.pythonhosted.org/packages/3a/13/547360d81e6d88d58492968ffda9f9542854f11310ee556fef14260cc886/zipp-4.1.0-py3-none-any.whl", hash = "sha256:25ad4e16390cd314347dd8f1de67a2ac538ae658ed4ab9db16029c07c188e97f", size = 10238, upload-time = "2026-05-18T20:08:57.045Z" },
-]

From cb49faf8ea8c73084fa6c7d44b9b7b95be7441c0 Mon Sep 17 00:00:00 2001
From: Evan Palmer <evanp922@gmail.com>
Date: Tue, 1 Sep 2026 13:57:51 -0700
Subject: [PATCH 3/4] Flatten envs/ package into common/env.py

The envs/ subpackage held only env.py after the MJX removal. Move it to
lrax/common/env.py and update the five relative imports (algorithm.py,
trainer.py, ppo.py, sac.py, shac.py) and the README example.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018a38TsYHunZKi4F8smG6JQ
---
 README.md                     |  2 +-
 lrax/common/algorithm.py      |  2 +-
 lrax/common/{envs => }/env.py |  2 +-
 lrax/common/envs/__init__.py  | 23 -----------------------
 lrax/common/trainer.py        |  2 +-
 lrax/ppo/ppo.py               |  2 +-
 lrax/sac/sac.py               |  2 +-
 lrax/shac/shac.py             |  2 +-
 8 files changed, 7 insertions(+), 30 deletions(-)
 rename lrax/common/{envs => }/env.py (98%)
 delete mode 100644 lrax/common/envs/__init__.py

diff --git a/README.md b/README.md
index 510657b..86fa858 100644
--- a/README.md
+++ b/README.md
@@ -43,7 +43,7 @@ import jax.random as jr
 import optax
 from jaxtyping import Array, PRNGKeyArray
 from lrax import PPO
-from lrax.common.envs import AbstractEnv, EnvState
+from lrax.common.env import AbstractEnv, EnvState
 from lrax.common.policies import Actor, ActorCritic, Critic
 from lrax.common.trainer import PolicyTrainer
 
diff --git a/lrax/common/algorithm.py b/lrax/common/algorithm.py
index f76ffa8..b47a96c 100644
--- a/lrax/common/algorithm.py
+++ b/lrax/common/algorithm.py
@@ -25,7 +25,7 @@
 from optax import OptState
 
 from .._custom_types import Metrics, Optimizers
-from .envs import AbstractEnv, EnvState
+from .env import AbstractEnv, EnvState
 
 type _AlgState = PyTree
 
diff --git a/lrax/common/envs/env.py b/lrax/common/env.py
similarity index 98%
rename from lrax/common/envs/env.py
rename to lrax/common/env.py
index f52f4f7..4019637 100644
--- a/lrax/common/envs/env.py
+++ b/lrax/common/env.py
@@ -23,7 +23,7 @@
 import equinox as eqx
 from jaxtyping import Array, PRNGKeyArray, PyTree
 
-from ..._custom_types import Metrics
+from .._custom_types import Metrics
 
 
 class EnvState(eqx.Module):
diff --git a/lrax/common/envs/__init__.py b/lrax/common/envs/__init__.py
deleted file mode 100644
index 02b6123..0000000
--- a/lrax/common/envs/__init__.py
+++ /dev/null
@@ -1,23 +0,0 @@
-# Copyright 2026, Laboratory for Robotics and Applied Mechanics
-#
-# Permission is hereby granted, free of charge, to any person obtaining a copy
-# of this software and associated documentation files (the "Software"), to deal
-# in the Software without restriction, including without limitation the rights
-# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
-# copies of the Software, and to permit persons to whom the Software is
-# furnished to do so, subject to the following conditions:
-#
-# The above copyright notice and this permission notice shall be included in
-# all copies or substantial portions of the Software.
-#
-# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
-# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
-# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL
-# THE AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
-# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
-# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
-# THE SOFTWARE.
-
-from .env import AbstractEnv, EnvState
-
-__all__ = ["AbstractEnv", "EnvState"]
diff --git a/lrax/common/trainer.py b/lrax/common/trainer.py
index 8132d65..740ad14 100644
--- a/lrax/common/trainer.py
+++ b/lrax/common/trainer.py
@@ -35,7 +35,7 @@
 from .._custom_types import Metrics, Optimizer
 from .algorithm import AbstractAlgorithm
 from .callbacks import Callback, StopTraining
-from .envs import AbstractEnv
+from .env import AbstractEnv
 from .logger import Logger
 from .policies import ActorCritic
 
diff --git a/lrax/ppo/ppo.py b/lrax/ppo/ppo.py
index b3159bd..b7918d6 100644
--- a/lrax/ppo/ppo.py
+++ b/lrax/ppo/ppo.py
@@ -32,7 +32,7 @@
 from .._custom_types import Metrics, Optimizer
 from .._epsilon import EPSILON
 from ..common.algorithm import AbstractAlgorithm
-from ..common.envs import AbstractEnv, EnvState
+from ..common.env import AbstractEnv, EnvState
 from ..common.policies import ActorCritic
 
 type _AlgState = None
diff --git a/lrax/sac/sac.py b/lrax/sac/sac.py
index 2f784fd..d6cd287 100644
--- a/lrax/sac/sac.py
+++ b/lrax/sac/sac.py
@@ -32,7 +32,7 @@
 from .._custom_types import Metrics, Optimizer
 from ..common.algorithm import AbstractAlgorithm
 from ..common.buffers import ReplayBuffer, Transition
-from ..common.envs import AbstractEnv, EnvState
+from ..common.env import AbstractEnv, EnvState
 from ..common.policies import ContinuousCritic
 from .policies import ActorCritic
 
diff --git a/lrax/shac/shac.py b/lrax/shac/shac.py
index 45a67fe..3aa2343 100644
--- a/lrax/shac/shac.py
+++ b/lrax/shac/shac.py
@@ -34,7 +34,7 @@
 from .._custom_types import Metrics, Optimizer
 from .._epsilon import EPSILON
 from ..common.algorithm import AbstractAlgorithm
-from ..common.envs import AbstractEnv, EnvState
+from ..common.env import AbstractEnv, EnvState
 from ..common.normalizers import RunningMeanStd
 from .policies import ActorCritic, Critic
 

From 0b3bcc6fc8d4e5689ac4f5990aaa890d7ecfba22 Mon Sep 17 00:00:00 2001
From: Evan Palmer <evanp922@gmail.com>
Date: Tue, 1 Sep 2026 14:00:13 -0700
Subject: [PATCH 4/4] Add issue and pull request templates

Mirrors the lab-standard templates from fiber / the rdml repos: a
PULL_REQUEST_TEMPLATE.md and bug-report / feature-request / documentation
issue forms with blank issues disabled. The bug report's environment
prompt is adapted to ask for lrax / Python / JAX versions and CPU vs GPU
rather than the ROS runtime details.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Claude-Session: https://claude.ai/code/session_018a38TsYHunZKi4F8smG6JQ
---
 .github/ISSUE_TEMPLATE/bug-report.yml      | 63 ++++++++++++++++++++++
 .github/ISSUE_TEMPLATE/config.yml          |  1 +
 .github/ISSUE_TEMPLATE/documentation.yml   | 46 ++++++++++++++++
 .github/ISSUE_TEMPLATE/feature-request.yml | 57 ++++++++++++++++++++
 .github/PULL_REQUEST_TEMPLATE.md           | 15 ++++++
 5 files changed, 182 insertions(+)
 create mode 100644 .github/ISSUE_TEMPLATE/bug-report.yml
 create mode 100644 .github/ISSUE_TEMPLATE/config.yml
 create mode 100644 .github/ISSUE_TEMPLATE/documentation.yml
 create mode 100644 .github/ISSUE_TEMPLATE/feature-request.yml
 create mode 100644 .github/PULL_REQUEST_TEMPLATE.md

diff --git a/.github/ISSUE_TEMPLATE/bug-report.yml b/.github/ISSUE_TEMPLATE/bug-report.yml
new file mode 100644
index 0000000..7436e8a
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/bug-report.yml
@@ -0,0 +1,63 @@
+name: Bug Report
+description: Report a bug.
+title: "<Please write a descriptive title>"
+labels: [needs triage]
+
+body:
+  - type: markdown
+    attributes:
+      value: >
+        Thank you for taking the time to file a bug report! Before creating a new
+        issue, please make sure to take a few minutes to check the issue tracker
+        for existing issues about the bug.
+
+  - type: textarea
+    attributes:
+      label: "Issue Description"
+      description: >
+        Please provide a clear and concise description of what the bug is.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Steps to Reproduce"
+      description: >
+        Please provide the steps that should be taken to reproduce the bug.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Expected Behavior"
+      description: >
+        Please describe or show an example of the expected behavior.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Error Message"
+      description: >
+        Please include the full error message, if any.
+      placeholder: >
+        << Full traceback starting from `Traceback: ...` >>
+      render: bash
+
+  - type: textarea
+    attributes:
+      label: "Runtime Environment"
+      description: >
+        Please provide the versions of lrax, Python, and JAX (jax and jaxlib),
+        the operating system, and whether JAX is running on CPU or GPU.
+      placeholder: >
+        lrax 0.1.0, Python 3.14, jax 0.10.2 / jaxlib 0.10.2 on CUDA GPU,
+        Ubuntu 24.04.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Additional Context"
+      description: >
+        Please provide any additional context needed to understand the bug.
diff --git a/.github/ISSUE_TEMPLATE/config.yml b/.github/ISSUE_TEMPLATE/config.yml
new file mode 100644
index 0000000..3ba13e0
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/config.yml
@@ -0,0 +1 @@
+blank_issues_enabled: false
diff --git a/.github/ISSUE_TEMPLATE/documentation.yml b/.github/ISSUE_TEMPLATE/documentation.yml
new file mode 100644
index 0000000..a7ce602
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/documentation.yml
@@ -0,0 +1,46 @@
+name: Documentation Improvement
+description: Report an issue related to the project documentation.
+title: "<Please write a descriptive title>"
+labels: [needs triage]
+
+body:
+  - type: markdown
+    attributes:
+      value: >
+        Thank you for taking the time to report a documentation issue! Before creating
+        a new issue, please make sure to take a few minutes to check the issue
+        tracker for existing issues similar to that being reported.
+
+  - type: dropdown
+    attributes:
+      label: Documentation Change Type
+      description: Please indicate what type of documentation issue you are reporting.
+      options:
+        - Adding new documentation to the project
+        - Changing existing project documentation
+        - Removing existing project documentation
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: Documentation Location
+      description: >
+        Please provide the location of the documentation that should be modified.
+
+  - type: textarea
+    attributes:
+      label: Documentation Problem
+      description: >
+        Please provide a description of how the documentation needs to be improved.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: Suggested Change
+      description: >
+        Please provide a description of the proposed change and why the proposed change
+        improves the upon the existing documentation.
+    validations:
+      required: true
diff --git a/.github/ISSUE_TEMPLATE/feature-request.yml b/.github/ISSUE_TEMPLATE/feature-request.yml
new file mode 100644
index 0000000..b69d8d1
--- /dev/null
+++ b/.github/ISSUE_TEMPLATE/feature-request.yml
@@ -0,0 +1,57 @@
+name: Feature Request
+description: Suggest a new idea for the project.
+title: "<Please write a descriptive title>"
+labels: [needs triage]
+
+body:
+  - type: markdown
+    attributes:
+      value: >
+        Thank you for taking the time to request a new feature! Before creating
+        a new issue, please make sure to take a few minutes to check the issue
+        tracker for existing issues similar to the proposed feature.
+
+  - type: dropdown
+    attributes:
+      label: Feature Type
+      description: Please indicate what type of feature request you would like to propose.
+      options:
+        - Adding new functionality to the project
+        - Changing existing functionality in the project
+        - Removing existing functionality in the project
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Problem Description"
+      description: >
+        Please provide a clear and concise description of what problem
+        the feature would solve.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Feature Description"
+      description: >
+        Please provide a description of the proposed feature, using pseudocode
+        if relevant.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Alternative Solutions"
+      description: >
+        Please provide a description of any alternative solutions or features
+        that would satisfy the feature request.
+    validations:
+      required: true
+
+  - type: textarea
+    attributes:
+      label: "Additional Context"
+      description: >
+        Please provide any additional context (e.g., relevant GitHub issues,
+        code examples, or references) needed to understand the feature request.
diff --git a/.github/PULL_REQUEST_TEMPLATE.md b/.github/PULL_REQUEST_TEMPLATE.md
new file mode 100644
index 0000000..8d1f140
--- /dev/null
+++ b/.github/PULL_REQUEST_TEMPLATE.md
@@ -0,0 +1,15 @@
+## Changes Made
+
+Please provide a description of all changes made in this PR and why the changes
+are needed.
+
+## Associated Issues
+
+Please provide a list of all open issues that this PR will close or contribute
+toward closing.
+
+- Fixes # (issue)
+
+## Testing
+
+Please provide a clear and concise description of the testing performed.