Add CI/CD pipelines

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

.github/ISSUE_TEMPLATE/config.yml

@@ -0,0 +1 @@
+blank_issues_enabled: false

.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

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

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

.github/dependabot.yml

@@ -0,0 +1,6 @@
+version: 2
+updates:
+  - package-ecosystem: github-actions
+    directory: "/"
+    schedule:
+      interval: "weekly"

.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

.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

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

.gitignore

@@ -1,2 +1,4 @@
 __pycache__/
 .venv/
+.ruff_cache/
+dist/

.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

README.md

@@ -25,10 +25,11 @@ uv pip install git+https://github.com/OSU-LRAM/lrax.git
 uv add git+ssh://git@github.com/OSU-LRAM/lrax.git
 ```
 
-To install the MuJoCo helpers, include the optional `mjx` dependencies
+The base install pulls in CPU-only JAX. For CUDA-accelerated JAX on Linux,
+include the optional `cuda` dependencies
 
 ```bash
-uv add "lrax[mjx] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
+uv add "lrax[cuda] @ git+ssh://git@github.com/OSU-LRAM/lrax.git"
 ```
 
 ## Usage
@@ -42,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
 
@@ -80,9 +81,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()

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