Kinematics and Dynamics

ambersim/utils/introspection_utils.py

@@ -15,6 +15,11 @@ def get_equality_names(model: mj.MjModel) -> List[str]:
     return [mj.mj_id2name(model, mj.mjtObj.mjOBJ_EQUALITY, i) for i in range(model.neq)]
 
 
+def get_body_names(model: mj.MjModel) -> List[str]:
+    """Returns a list of all geom names in a mujoco (NOT mjx) model."""
+    return [mj.mj_id2name(model, mj.mjtObj.mjOBJ_BODY, i) for i in range(model.nbody)]
+
+
 def get_geom_names(model: mj.MjModel) -> List[str]:
     """Returns a list of all geom names in a mujoco (NOT mjx) model."""
     return [mj.mj_id2name(model, mj.mjtObj.mjOBJ_GEOM, i) for i in range(model.ngeom)]

ambersim/utils/io_utils.py

@@ -2,17 +2,17 @@
 from typing import Optional, Tuple, Union
 
 import coacd
-import mujoco as mj
 import numpy as np
 import trimesh
 from dm_control import mjcf
 from lxml import etree
-from mujoco import mjx
 from packaging import version
 
+import mujoco as mj
 from ambersim import ROOT
 from ambersim.utils._internal_utils import _check_filepath
 from ambersim.utils.conversion_utils import save_model_xml
+from mujoco import mjx
 
 
 def _add_actuators(urdf_filepath: Union[str, Path], xml_filepath: Union[str, Path]) -> None:

ambersim/utils/pin_utils.py

@@ -0,0 +1,159 @@
+from typing import List, Tuple
+
+import jax
+import mujoco as mj
+import numpy as np
+import pinocchio as pin
+from mujoco import mjx
+from mujoco.mjx._src.scan import _q_jointid
+
+from ambersim.utils.introspection_utils import get_joint_names
+
+"""Utils for converting quantities between mujoco/mjx and pinocchio."""
+
+
+def get_joint_orders(
+    mj_model: mj.MjModel,
+    pin_model: pin.Model,
+) -> Tuple[List[int], List[int]]:
+    """Computes the mapping from joint orders in mujoco and pinocchio.
+
+    Pinocchio maintains one extra joint relative to mujoco (the "universe" joint).
+    The outputs satisfy the following two tests:
+
+    assert [mj_joint_names[i] if i is not None else "universe" for i in mj2pin] == pin_joint_names
+    assert [pin_joint_names[i] for i in pin2mj if i is not None] == mj_joint_names
+
+    Args:
+        mj_model: The MuJoCo model.
+        pin_model: The Pinocchio model.
+
+    Returns:
+        pin2mj: The mapping from pinocchio joint order to mujoco joint order.
+        mj2pin: The mapping from mujoco joint order to pinocchio joint order.
+    """
+    mj_joint_names = get_joint_names(mj_model)
+    pin_joint_names = list(pin_model.names)
+
+    mj2pin = [mj_joint_names.index(a) if a in mj_joint_names else None for a in pin_joint_names]
+    pin2mj = [pin_joint_names.index(a) if a in pin_joint_names else None for a in mj_joint_names]
+
+    return mj2pin, pin2mj
+
+
+def mjx_xanchor_to_pin(mj2pin: List[int], mjx_data: mjx.Data) -> np.ndarray:
+    """Converts the mujoco joint anchor positions to pinocchio joint placements.
+
+    Args:
+        mj2pin: The mapping from mujoco joint order to pinocchio joint order.
+        mjx_data: The mujoco data.
+
+    Returns:
+        xanchor_pin: The pinocchio joint placements.
+    """
+    xanchors = np.array(mjx_data.xanchor)
+    _xanchor_pin = []
+    for mj_idx in mj2pin:
+        if mj_idx is not None:
+            _xanchor_pin.append(xanchors[mj_idx, :])
+        else:
+            _xanchor_pin.append(np.zeros(3))
+    return np.stack(_xanchor_pin)
+
+
+def pin_xanchor_to_mjx(pin2mj: List[int], pin_data: pin.Data) -> np.ndarray:
+    """Converts the pinocchio joint placements to mujoco joint anchor positions.
+
+    Args:
+        pin2mj: The mapping from pinocchio joint order to mujoco joint order.
+        pin_data: The pinocchio data.
+
+    Returns:
+        xanchor_mj: The mujoco joint anchor positions.
+    """
+    _xanchor_mj = []
+    for pin_idx in pin2mj:
+        if pin_idx is not None:
+            _xanchor_mj.append(pin_data.oMi[pin_idx].translation)
+        else:
+            _xanchor_mj.append(np.zeros(3))
+    return np.stack(_xanchor_mj)
+
+
+def mjx_qpos_to_pin(qpos_mjx: jax.Array, mjx_model: mjx.Model, mj2pin: List[int]) -> np.ndarray:
+    """Converts the mujoco joint positions to pinocchio joint positions.
+
+    Args:
+        qpos_mjx: The mjx joint positions.
+        mjx_model: The mjx model.
+        mj2pin: The mapping from mujoco joint order to pinocchio joint order.
+
+    Returns:
+        qpos_pin: The pinocchio joint positions.
+    """
+    qpos_mjx = np.array(qpos_mjx)
+    _qpos_pin = []
+    jnt_types = mjx_model.jnt_type
+    for mj_idx in mj2pin:
+        # ignore Nones
+        if mj_idx is None:
+            continue
+
+        # check how to append based on joint type
+        jnt_type = jnt_types[mj_idx]
+        startidx = mjx_model.jnt_qposadr[mj_idx]
+
+        if jnt_type == 0:  # FREE, 7-dimensional
+            # pinocchio uses the (x, y, z, w) convention
+            _qpos_pin.append(qpos_mjx[startidx : startidx + 3])
+            quat_wxyz = qpos_mjx[startidx + 3 : startidx + 7]
+            quat_xyzw = np.concatenate((quat_wxyz[1:], quat_wxyz[:1]))
+            _qpos_pin.append(quat_xyzw)
+        elif jnt_type == 1:  # BALL, 4-dimensional
+            quat_wxyz = qpos_mjx[startidx : startidx + 4]
+            quat_xyzw = np.concatenate((quat_wxyz[1:], quat_wxyz[:1]))
+            _qpos_pin.append(quat_xyzw)
+        elif jnt_type == 2:  # SLIDE, 3-dimensional
+            _qpos_pin.append(qpos_mjx[startidx : startidx + 3])
+        else:  # HINGE, 1-dimensional
+            _qpos_pin.append(qpos_mjx[startidx : startidx + 1])
+
+    return np.concatenate(_qpos_pin)
+
+
+def mjx_qvel_to_pin(qvel_mjx: jax.Array, mjx_model: mjx.Model, mj2pin: List[int]) -> np.ndarray:
+    """Converts the mujoco joint velocities to pinocchio joint velocities.
+
+    Notice that there is one less coordinate for rotational velocities compared to quaternion coordinates.
+
+    Args:
+        qvel_mjx: The mjx joint velocities.
+        mjx_model: The mjx model.
+        mj2pin: The mapping from mujoco joint order to pinocchio joint order.
+
+    Returns:
+        qvel_pin: The pinocchio joint velocities.
+    """
+    qvel_mjx = np.array(qvel_mjx)
+    _qvel_pin = []
+    jnt_types = mjx_model.jnt_type
+    for mj_idx in mj2pin:
+        # ignore Nones
+        if mj_idx is None:
+            continue
+
+        # check how to append based on joint type
+        jnt_type = jnt_types[mj_idx]
+        startidx = mjx_model.jnt_dofadr[mj_idx]
+
+        if jnt_type == 0:  # FREE, 6-dimensional
+            _qvel_pin.append(qvel_mjx[startidx : startidx + 3])
+            _qvel_pin.append(qvel_mjx[startidx + 3 : startidx + 6])
+        elif jnt_type == 1:  # BALL, 3-dimensional
+            _qvel_pin.append(qvel_mjx[startidx : startidx + 3])
+        elif jnt_type == 2:  # SLIDE, 1-dimensional
+            _qvel_pin.append(qvel_mjx[startidx : startidx + 1])
+        else:  # HINGE, 1-dimensional
+            _qvel_pin.append(qvel_mjx[startidx : startidx + 1])
+
+    return np.concatenate(_qvel_pin)

install.sh

@@ -0,0 +1,36 @@
+#!/bin/bash
+# instructions for installing from source are taken from two places:
+# (1) https://mujoco.readthedocs.io/en/latest/programming/#building-from-source
+# (2) https://mujoco.readthedocs.io/en/latest/python.html#building-from-source
+
+set -e
+
+dev=false
+source=false
+while getopts dsh: flag; do
+    case "${flag}" in
+        d) dev=true;;   # Install development dependencies
+        s) source=true;; # Install MuJoCo from source
+        h) hash=${OPTARG};;  # Hash of the MuJoCo commit to install
+    esac
+done
+
+# Install regular or development dependencies
+if [ "$dev" = true ] ; then
+    echo "[NOTE] Installing development dependencies..."
+    pip install --upgrade --upgrade-strategy only-if-needed -e .[all] \
+        --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+        --find-links https://download.pytorch.org/whl/cu118
+    pre-commit autoupdate
+    pre-commit install
+else
+    echo "[NOTE] Installing non-developer dependencies..."
+    pip install --upgrade --upgrade-strategy only-if-needed -e . --default-timeout=100 future \
+        --find-links https://storage.googleapis.com/jax-releases/jax_cuda_releases.html \
+        --find-links https://download.pytorch.org/whl/cu118
+fi
+
+# Checking whether to install mujoco from source
+if [ "$source" = true ] ; then
+    bash install_mj_source.sh -h "$hash"
+fi

install_mj_source.sh

@@ -0,0 +1,101 @@
+#!/bin/bash
+
+while getopts h: flag; do
+    case "${flag}" in
+        h) hash=${OPTARG};;  # Hash of the MuJoCo commit to install
+    esac
+done
+
+echo -e "\n[NOTE] Installing mujoco from source..."
+
+# the script directory and mujoco directory
+script_dir="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+mujoco_dir="$script_dir/../mujoco"
+
+# check whether we already have the most recent release cached to save build time
+LATEST_GIT_HASH=$(git ls-remote https://github.com/google-deepmind/mujoco HEAD | awk '{ print $1}')
+if [ -d "$mujoco_dir/python/dist" ]; then
+    tar_path=$(find "$mujoco_dir/python/dist" -name "mujoco-*${LATEST_GIT_HASH}.tar.gz" 2>/dev/null)
+    whl_path=$(find "$mujoco_dir/python/dist" -name "mujoco-*.whl" 2>/dev/null)
+
+    if [ -f "$tar_path" ]; then
+        echo -e "[NOTE] Found cached mujoco tar.gz file..."
+        if [ -f "$whl_path" ]; then
+            echo -e "[NOTE] Wheel found! Installing from wheel..."
+            cd "$mujoco_dir/python/dist"
+
+            # [Nov 10, 2023] we install with --no-deps because this would upgrade numpy to a version incompatible with cmeel-boost 1.82.0
+            MUJOCO_PATH="$mujoco_dir/mujoco_install" MUJOCO_PLUGIN_PATH="$mujoco_dir/plugin" pip install --no-deps --force-reinstall "$whl_path"
+            cd "$mujoco_dir/mjx"
+            pip install --no-deps --force-reinstall .
+            exit 0
+        else
+            echo -e "[NOTE] No wheel found! Building it and installing..."
+            cd "$mujoco_dir/python/dist"
+            MUJOCO_PATH="$mujoco_dir/mujoco_install" MUJOCO_PLUGIN_PATH="$mujoco_dir/plugin" pip wheel -w $(dirname "whl_path") "$tar_path"
+
+            # [Nov 10, 2023] we install with --no-deps because this would upgrade numpy to a version incompatible with cmeel-boost 1.82.0
+            MUJOCO_PATH="$mujoco_dir/mujoco_install" MUJOCO_PLUGIN_PATH="$mujoco_dir/plugin" pip install --no-deps --force-reinstall "$whl_path"
+            cd "$mujoco_dir/mjx"
+            pip install --no-deps --force-reinstall .
+            exit 0
+        fi
+    fi
+fi
+
+# Check if CMake and a compiler are installed
+command -v cmake &> /dev/null || { echo "CMake is not installed. Please install CMake before proceeding."; exit 1; }
+command -v g++ &> /dev/null || { echo "C++ compiler is not available. Please install a C++ compiler before proceeding."; exit 1; }
+
+
+# Clone the Mujoco repo to the directory above where this script is located
+# If it exists already, then just git pull new changes
+if [ -d "$mujoco_dir" ]; then
+    echo "Mujoco exists already - running git pull to update it!"
+    cd "$mujoco_dir"
+    git pull origin main
+    if [ ! -z "$hash" ]; then
+        echo "Checking out with provided commit hash!"
+        git checkout "$hash"
+    fi
+else
+    echo "Mujoco does not exist - cloning it!"
+    git clone https://github.com/google-deepmind/mujoco.git "$mujoco_dir"
+    if [ ! -z "$hash" ]; then
+        echo "Checking out to provided hash!"
+        (cd "$mujoco_dir" && git checkout "$hash")
+    fi
+fi
+
+# Configure and build
+export MAKEFLAGS="-j$(nproc)"  # allow the max number of processors when building
+cd "$mujoco_dir"
+if [ -z "$hash" ]; then
+    SAVED_GIT_HASH=$(git rev-parse HEAD)  # getting the git hash
+else
+    SAVED_GIT_HASH="$hash"
+fi
+cmake . && cmake --build .
+
+# Install Mujoco
+cmake . -DCMAKE_INSTALL_PREFIX="./mujoco_install" && cmake --install .
+
+# Generate source distribution required for Python bindings
+cd "$mujoco_dir/python"
+./make_sdist.sh
+tar_path=$(find "$mujoco_dir/python/dist" -name 'mujoco-*.tar.gz' 2>/dev/null)
+
+# Renaming the tar file by appending the commit hash
+new_tar_path="$(dirname "$tar_path")/$(basename "$tar_path" .tar.gz)-$SAVED_GIT_HASH.tar.gz"
+mv "$tar_path" "$new_tar_path"
+
+# manually building wheel so we can cache it
+cd "$mujoco_dir/python/dist"
+MUJOCO_PATH="$mujoco_dir/mujoco_install" MUJOCO_PLUGIN_PATH="$mujoco_dir/plugin" pip wheel -w $(dirname "new_tar_path") "$new_tar_path"
+
+# installing mujoco from wheel and then finally mjx
+whl_path=$(find "$mujoco_dir/python/dist" -name "mujoco-*.whl" 2>/dev/null)
+MUJOCO_PATH="$mujoco_dir/mujoco_install" MUJOCO_PLUGIN_PATH="$mujoco_dir/plugin" pip install --no-deps --force-reinstall "$whl_path"
+cd "$mujoco_dir/mjx"
+pip wheel -w "$mujoco_dir/mjx" --no-deps .
+pip install --no-deps --force-reinstall .