Allow Installation of mujoco from Source

.github/workflows/code_checks.yml

@@ -14,30 +14,71 @@ jobs:
     runs-on: ubuntu-latest
 
     steps:
+      # checks out ambersim
       - uses: actions/checkout@v4
-      - name: Set up Python 3.11.5
-        uses: actions/setup-python@v4
-        with:
-          python-version: '3.11.5'
+
+      # checks the cache for pip packages
       - uses: actions/cache@v3
         with:
           path: ~/.cache/pip
           key: ${{ runner.os }}-${{ hashFiles('pyproject.toml') }}
           restore-keys: |
             ${{ runner.os }}-pip-
+
+      # download the cached artifacts
+      - name: Download mujoco artifact
+        id: download-artifact-mujoco
+        uses: dawidd6/action-download-artifact@v2
+        with:
+          workflow: mujoco_nightly.yml
+          name_is_regexp: true
+          name: "mujoco_wheel-*"
+          path: ${{ github.workspace }}
+          check_artifacts: true
+          search_artifacts: true
+          if_no_artifact_found: warn
+      - name: Download mjx artifact
+        id: download-artifact-mjx
+        uses: dawidd6/action-download-artifact@v2
+        with:
+          workflow: mujoco_nightly.yml
+          name_is_regexp: true
+          name: "mjx_wheel-*"
+          path: ${{ github.workspace }}
+          check_artifacts: true
+          search_artifacts: true
+          if_no_artifact_found: warn
       - name: Install dependencies
+        working-directory: ${{ github.workspace }}
+        shell: bash -l {0}
         run: |
           python -m pip install --upgrade pip
-          pip install --upgrade --upgrade-strategy eager -e .[all] --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --find-links https://download.pytorch.org/whl/cu118
+          bash ./install.sh -d  # -d means dev, don't use -s because we manually install mj/mjx from source here
+
+          # upgrade mujoco to use nightly build
+          mujoco_whl_path=$(find ${{ github.workspace }} -name "mujoco-*.whl")
+          mjx_whl_path=$(find ${{ github.workspace }} -name "mujoco_mjx-*.whl")
+          if [ -n "$mujoco_whl_path" ]; then
+              pip install --no-deps --force-reinstall $mujoco_whl_path
+          fi
+          if [ -n "$mjx_whl_path" ]; then
+              pip install --no-deps --force-reinstall $mjx_whl_path
+          fi
+
+      # run all code checks
       - name: Run black
+        shell: bash -l {0}
         run: black --check .
       - name: Run flake8
+        shell: bash -l {0}
         run: |
           # stop the build if there are Python syntax errors or undefined names
           flake8 . --count --select=E9,F63,F7,F82 --show-source --statistics
       - name: Run pyright
+        shell: bash -l {0}
         run: |
           pyright
       - name: Test with pytest
+        shell: bash -l {0}
         run: |
           pytest

.github/workflows/mujoco_nightly.yml

@@ -0,0 +1,157 @@
+name: Build Nightly MuJoCo Release
+
+on:
+  schedule:
+    - cron: "53 12 * * *"  # runs at 4:53AM PST every day - avoid high load on the hour
+  pull_request:
+    branches: [main]
+    types: [ready_for_review]  # also caches when a PR is ready for review
+
+jobs:
+  build-and-package:
+    runs-on: ubuntu-latest
+
+    steps:
+      # checks out ambersim
+      - name: Checkout code
+        uses: actions/checkout@v4
+
+      # frees up space in the runner, since there's not enough disk space to build mujoco from source
+      - name: Free Up GitHub Actions Ubuntu Runner Disk Space 🔧
+        uses: jlumbroso/free-disk-space@main
+        with:
+          tool-cache: false  # might have useful/necessary stuff
+          android: true  # removes 14GB very quickly
+          dotnet: false  # removes a lot of space, but unnecessary here. skip to save time.
+          haskell: false  # removes a lot of space, but unnecessary here. skip to save time.
+          large-packages: false  # doesn't save that much space relatively, takes long to run
+          swap-storage: false  # removes a lot of space, but unnecessary here. skip to save time.
+
+      # for caching the conda env - re-updates the env every 24 hours just in case
+      # see: https://github.com/conda-incubator/setup-miniconda#caching-environments
+      - name: Setup Mambaforge
+        uses: conda-incubator/setup-miniconda@v2
+        with:
+          miniforge-variant: Mambaforge
+          miniforge-version: latest
+          activate-environment: anaconda-client-env
+          use-mamba: true
+      - name: Get Date
+        id: get-date
+        run: echo "today=$(/bin/date -u '+%Y%m%d')" >> $GITHUB_OUTPUT
+        shell: bash
+      - name: Cache conda
+        uses: actions/cache@v3
+        env:
+          CACHE_NUMBER: 0
+        with:
+          path: ${{ env.CONDA }}/envs
+          key:
+            conda-${{ runner.os }}--${{ runner.arch }}--${{
+            steps.get-date.outputs.today }}-${{
+            hashFiles('etc/example-environment-caching.yml') }}-${{ env.CACHE_NUMBER
+            }}
+        id: cache
+      - name: Update environment
+        run:
+          mamba env update -n ambersim -f environment.yml
+        if: steps.cache.outputs.cache-hit != 'true'  # only re-initializes the environment if we don't get a cache hit
+
+      # download previously built mujoco and mjx artifacts
+      - name: Download mujoco artifact
+        id: download-artifact-mujoco
+        uses: dawidd6/action-download-artifact@v2
+        with:
+          workflow: mujoco_nightly.yml
+          name_is_regexp: true
+          name: "mujoco_wheel-*"
+          path: ${{ github.workspace }}
+          check_artifacts: true
+          search_artifacts: true
+          if_no_artifact_found: warn
+      - name: Download mjx artifact
+        id: download-artifact-mjx
+        uses: dawidd6/action-download-artifact@v2
+        with:
+          workflow: mujoco_nightly.yml
+          name_is_regexp: true
+          name: "mjx_wheel-*"
+          path: ${{ github.workspace }}
+          check_artifacts: true
+          search_artifacts: true
+          if_no_artifact_found: warn
+
+      # checks whether to rebuild or not
+      - name: Build pre-check
+        shell: bash -l {0}
+        run: |
+          # Checks the git hash of the cached files vs. the most recent one
+          LATEST_GIT_HASH=$(git ls-remote https://github.com/google-deepmind/mujoco HEAD | awk '{ print $1}')
+
+          mujoco_whl_path=$(find ${{ github.workspace }} -name "mujoco-*.whl")
+          mjx_whl_path=$(find ${{ github.workspace }} -name "mujoco_mjx-*.whl")
+
+          if [ -n "$mujoco_whl_path" ]; then
+              # if the file name exists, extract the git hash and compare to upstream mujoco
+              CACHED_GIT_HASH=$(echo "$mujoco_whl_path" | grep -oP "mujoco_wheel-\K[^/]+")
+              if [ "$LATEST_GIT_HASH" = "$CACHED_GIT_HASH" ]; then
+                echo "Cached wheel matches most recent one - skipping build! Hash: $CACHED_GIT_HASH"
+                echo "skip_build=true" >> $GITHUB_ENV
+
+                # save the names and filepaths so we can just re-upload them
+                echo "mujoco_name=mujoco_wheel-${LATEST_GIT_HASH}" >> $GITHUB_ENV
+                echo "mjx_name=mjx_wheel-${LATEST_GIT_HASH}" >> $GITHUB_ENV
+                echo "mujoco_path=$(realpath $mujoco_whl_path)" >> $GITHUB_ENV
+                echo "mjx_path=$(realpath $mjx_whl_path)" >> $GITHUB_ENV
+              else
+                echo "Cached wheel is outdated - building new ones! Old Hash: $CACHED_GIT_HASH"
+                echo "skip_build=false" >> $GITHUB_ENV
+              fi
+          else
+              echo "No cached wheels found - building new ones!"
+              echo "skip_build=false" >> $GITHUB_ENV
+          fi
+
+      # mujoco build + install
+      - name: Create and Store Wheel
+        shell: bash -l {0}
+        run: |
+          # Installs mujoco from source
+          sudo apt-get update -y
+          sudo apt-get install -y \
+            libgl1-mesa-dev \
+            libxinerama-dev \
+            libxcursor-dev \
+            libxrandr-dev \
+            libxi-dev \
+            ninja-build
+          python -m pip install --upgrade pip
+          ./install_mj_source.sh
+
+          # Saving the wheels to respective paths
+          mujoco_whl_path=$(find $PWD/../mujoco/python/dist -name "mujoco-*.whl")
+          mjx_whl_path=$(find $PWD/../mujoco/mjx -name "mujoco_mjx-*.whl")
+          echo "mujoco_path=$(realpath $mujoco_whl_path)" >> $GITHUB_ENV
+          echo "mjx_path=$(realpath $mjx_whl_path)" >> $GITHUB_ENV
+
+          # Appending the relevant git commit hash to the artifact name
+          LATEST_GIT_HASH=$(git ls-remote https://github.com/google-deepmind/mujoco HEAD | awk '{ print $1}')
+          echo "mujoco_name=mujoco_wheel-${LATEST_GIT_HASH}" >> $GITHUB_ENV
+          echo "mjx_name=mjx_wheel-${LATEST_GIT_HASH}" >> $GITHUB_ENV
+        id: package
+        if: env.skip_build != 'true'
+
+      # upload the built wheels as artifacts
+      - name: Upload mujoco artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: ${{ env.mujoco_name }}
+          path: ${{ env.mujoco_path }}
+          retention-days: 7
+
+      - name: Upload mjx artifact
+        uses: actions/upload-artifact@v3
+        with:
+          name: ${{ env.mjx_name }}
+          path: ${{ env.mjx_path }}
+          retention-days: 7

.pre-commit-config.yaml

@@ -15,7 +15,7 @@ repos:
 
 # code format according to black
 -   repo: https://github.com/ambv/black
-    rev: 23.10.1
+    rev: 23.11.0
     hooks:
     -   id: black
 

README.md

@@ -4,33 +4,49 @@ This repository houses tools built on the GPU-accelerated simulation capabilitie
 * shared interfaces for control architectures, and
 * massively-parallelized simulation.
 
-## Quickstart
-
-### Non-developers
-Create a conda environment with Cuda 11.8 support:
+## Installation
+Clone this repository and run the following commands in the repository root to create and activate a conda environment with Cuda 11.8 support:
 ```
 conda env create -n <env_name> -f environment.yml
 conda activate <env_name>
 ```
-To locally install this package, clone the repository and in the repo root, run
+
+TL;DR: installation commands are here. This will ask you for your password to install system-wide dependencies. For details, see below.
+
+For non-developers installing `mujoco` from source:
 ```
-pip install . --default-timeout=100 future --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --find-links https://download.pytorch.org/whl/cu118
+./install.sh -s
 ```
 
-### Developers
-Create a conda environment with Cuda 11.8 support:
+For developers installing `mujoco` from source:
 ```
-conda env create -n <env_name> -f environment.yml
-conda activate <env_name>
+# no path to the mujoco repo specified
+./install.sh -s -d
+
+# specifying a path to the mujoco repo
+./install.sh -s -d --mujoco-dir /path/ending/in/mujoco
 ```
-Install the project with the editable flag and development dependencies:
+
+Installation of this package is done via the above `bash` script. There are a few flags for configuring the installation:
+* `-d` controls whether to use the heavier _development_ dependencies, which include linting and testing dependencies;
+* `-s` controls whether to install the most recent `mujoco` version from source. We recommend doing this, since the development version usually has important bugfixes.
+* `--disable-apt` specifies whether to disable the system-wide dependencies installed by `apt` which are required to install `mujoco` from source. They are enabled by default. The packages are:
+    * `libgl1-mesa-dev`
+    * `libxinerama-dev`
+    * `libxcursor-dev`
+    * `libxrandr-dev`
+    * `libxi-dev`
+    * `ninja-build`
+* `--mujoco-dir` specifies the directory of the local `mujoco` repo, which must end in the directory `mujoco`. If one doesn't exist, it will be pulled to this directory. If this isn't specified, `mujoco` will be created as a sibling directory of `ambersim`.
+
+If the following line of code runs without error, then the installation of `mujoco` from source was successful:
 ```
-pip install -e .[all] --default-timeout=100 future --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html --find-links https://download.pytorch.org/whl/cu118
+python -c "import mujoco; from mujoco import mjx"
 ```
-Then, install pre-commit hooks by running the following in the repo root:
+Further, you can examine the latest minor version using `pip`:
 ```
-pre-commit autoupdate
-pre-commit install
+pip show mujoco
+pip show mujoco-mjx
 ```
 
 ## Custom Models
@@ -45,6 +61,7 @@ We have implemented some custom utils for model parsing, but they aren't complet
 ### Abridged Dev Guidelines
 Development on this code will be controlled via code review. To facilitate this, please follow these guidelines:
 * keep your pull requests small so that it's practical to human review them;
+* try to create _draft pull requests_ instead of regular ones and request reviews from relevant people only when ready - we rebuild `mujoco` from source when this happens;
 * write tests as you go (and if you are reviewing, suggest missing tests);
 * write docstrings for public classes and methods, even if it's just a one-liner;
 * before committing, make sure you locally pass all tests by running `pytest` in the repo root;
@@ -56,7 +73,7 @@ Python dependencies are specified using a `pyproject.toml` file. Non-python depe
 
 Major versioning decisions:
 * `python=3.11.5`. `torch`, `jax`, and `mujoco` all support it and there are major reported speed improvements over `python` 3.10.
-* `cuda==11.8`. Both `torch` and `jax` support `cuda12`; however, they annoyingly support different minor versions which makes them incompatible in the same environment [https://github.com/google/jax/issues/18032](#18032). Once this is resolved, we will upgrade to `cuda-12.2` or later. It seems most likely that `torch` will support `cuda-12.3` once they do upgrade, since that is the most recent release.
+* `cuda==11.8`. Both `torch` and `jax` support `cuda12`; however, they annoyingly support different minor versions which makes them [incompatible in the same environment](https://github.com/google/jax/issues/18032). Once this is resolved, we will upgrade to `cuda-12.2` or later. It seems most likely that `torch` will support `cuda-12.3` once they do upgrade, since that is the most recent release.
 
 ### Tooling
 We use various tools to ensure code quality.

ambersim/utils/io_utils.py

@@ -8,6 +8,7 @@
 from dm_control import mjcf
 from lxml import etree
 from mujoco import mjx
+from packaging import version
 
 from ambersim import ROOT
 from ambersim.utils._internal_utils import _check_filepath
@@ -138,7 +139,7 @@ def _modify_robot_float_base(filepath: Union[str, Path]) -> mj.MjModel:
 def load_mj_model_from_file(
     filepath: Union[str, Path],
     force_float: bool = False,
-    solver: Union[str, mj.mjtSolver] = mj.mjtSolver.mjSOL_CG,
+    solver: Optional[Union[str, mj.mjtSolver]] = None,
     iterations: Optional[int] = None,
     ls_iterations: Optional[int] = None,
 ) -> mj.MjModel:
@@ -157,14 +158,29 @@ def load_mj_model_from_file(
     Raises:
         NotImplementedError: if the file extension is not in [".urdf", ".xml"]
     """
-    # TODO(ahl): once we allow installing mujoco from source, update this to allow the Newton solver + update default
-    if isinstance(solver, str):
-        if solver.lower() == "cg":
-            solver = mj.mjtSolver.SOL_CG
-        else:
-            raise ValueError("Solver must be one of: ['cg']!")
-    elif isinstance(solver, mj.mjtSolver):
-        assert solver in [mj.mjtSolver.mjSOL_CG]
+    # allow different solver specifications depending on mujoco version
+    if version.parse(mj.__version__) < version.parse("3.0.1"):
+        if solver is None:
+            solver = mj.mjtSolver.mjSOL_CG
+        elif isinstance(solver, str):
+            if solver.lower() == "cg":
+                solver = mj.mjtSolver.mjSOL_CG
+            else:
+                raise ValueError("Solver must be one of: ['cg']!")
+        elif isinstance(solver, mj.mjtSolver):
+            assert solver in [mj.mjtSolver.mjSOL_CG]
+    else:
+        if solver is None:
+            solver = mj.mjtSolver.mjSOL_NEWTON
+        elif isinstance(solver, str):
+            if solver.lower() == "newton":
+                solver = mj.mjtSolver.mjSOL_NEWTON
+            elif solver.lower() == "cg":
+                solver = mj.mjtSolver.mjSOL_CG
+            else:
+                raise ValueError("Solver must be one of: ['cg', 'newton']!")
+        elif isinstance(solver, mj.mjtSolver):
+            assert solver in [mj.mjtSolver.mjSOL_CG, mj.mjtSolver.mjSOL_NEWTON]
 
     filepath = _check_filepath(filepath)
     is_urdf = str(filepath).split(".")[-1] == "urdf"
@@ -176,7 +192,7 @@ def load_mj_model_from_file(
         output_path = "/".join(str(filepath).split("/")[:-1]) + "/_temp_xml_model.xml"
         save_model_xml(filepath, output_path=output_path)
         _add_actuators(filepath, output_path)
-        # _add_mimics(filepath, output_path)  # TODO(ahl): uncomment when we merge #21
+        _add_mimics(filepath, output_path)
         temp_output_path = True
     elif is_xml:
         output_path = filepath
@@ -191,7 +207,7 @@ def load_mj_model_from_file(
 
     # deleting temp file
     if temp_output_path:
-        Path.unlink(output_path)
+        Path(output_path).unlink()
 
     # setting solver options
     mj_model.opt.solver = solver