diff --git a/ambersim/utils/introspection_utils.py b/ambersim/utils/introspection_utils.py index 4183b856..9a03adfc 100644 --- a/ambersim/utils/introspection_utils.py +++ b/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)] diff --git a/ambersim/utils/io_utils.py b/ambersim/utils/io_utils.py index fe399f62..71106ba5 100644 --- a/ambersim/utils/io_utils.py +++ b/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: diff --git a/ambersim/utils/pin_utils.py b/ambersim/utils/pin_utils.py new file mode 100644 index 00000000..fb1a5a57 --- /dev/null +++ b/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) diff --git a/install.sh b/install.sh new file mode 100755 index 00000000..d8b1cd47 --- /dev/null +++ b/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 diff --git a/install_mj_source.sh b/install_mj_source.sh new file mode 100755 index 00000000..85f8a987 --- /dev/null +++ b/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 . \ No newline at end of file diff --git a/tests/test_kin_dyn.py b/tests/test_kin_dyn.py new file mode 100644 index 00000000..02c8ddad --- /dev/null +++ b/tests/test_kin_dyn.py @@ -0,0 +1,163 @@ +import jax.numpy as jnp +import mujoco as mj +import numpy as np +import pinocchio as pin +import pytest +from mujoco import mjx + +from ambersim import ROOT +from ambersim.utils.introspection_utils import get_joint_names +from ambersim.utils.io_utils import load_mj_model_from_file, mj_to_mjx_model_and_data +from ambersim.utils.pin_utils import ( + get_joint_orders, + mjx_qpos_to_pin, + mjx_qvel_to_pin, + mjx_xanchor_to_pin, + pin_xanchor_to_mjx, +) + +""" +The tests here check the mjx calculations against pinocchio (and also demonstrate how to use the kin/dyn funcs +independently of environments). The two main bookkeeping things are (1) the joint names and (2) the different +quaternion convention. + +TODO(ahl): +* test correctness of batched data +* use more convincing nonzero values for dynamics tests +""" + + +@pytest.fixture +def shared_variables(): + """Configuration method for global variables.""" + # run all tests on the barrett hand + urdf_path = ROOT + "/models/barrett_hand/bh280.urdf" + + # loading mjx models + mj_model = load_mj_model_from_file(urdf_path, force_float=True) + mjx_model, mjx_data = mj_to_mjx_model_and_data(mj_model) + mj_joint_names = get_joint_names(mj_model) + + # loading pinocchio models + pin_model = pin.buildModelFromUrdf(urdf_path, pin.JointModelFreeFlyer()) + pin_model.names[1] = "freejoint" # by default, adding a joint above names it "root_joint", so we rename it + pin_data = pin_model.createData() + pin_joint_names = list(pin_model.names) + + # relating mujoco and pinocchio joint orders + mj2pin, pin2mj = get_joint_orders(mj_model, pin_model) + + # setting dummy positions + _qpos_mj = [] + for mjjn in mj_joint_names: + if mjjn != "freejoint": + _qpos_mj.append([np.random.randn()]) # arbitrary + else: + # floating base state convention: (translation[x, y, z], quat[w, x, y, z]) + _translation = np.random.randn(3) + _uvw = np.random.rand(3) + _quat = np.array( + [ + np.sqrt(1.0 - _uvw[0]) * np.sin(2 * np.pi * _uvw[1]), + np.sqrt(1.0 - _uvw[0]) * np.cos(2 * np.pi * _uvw[1]), + np.sqrt(_uvw[0]) * np.sin(2 * np.pi * _uvw[2]), + np.sqrt(_uvw[0]) * np.cos(2 * np.pi * _uvw[2]), + ] + ) # uniformly random quaternion: stackoverflow.com/a/44031492 + _qpos_mj.append(np.concatenate((_translation, _quat))) + qpos_mjx = jnp.array(np.concatenate(_qpos_mj)) + qpos_pin = mjx_qpos_to_pin(qpos_mjx, mjx_model, mj2pin) # accounts for potentially different joint ordering + + # setting dummy velocities + # qvel_mjx = jnp.array(np.random.randn(mjx_model.nv)) + # qvel_pin = mjx_qvel_to_pin(qvel_mjx, mjx_model, mj2pin) # accounts for potentially different joint ordering + + # # setting dummy torques + # qtau_mjx = jnp.array(np.random.randn(mjx_model.nu)) + # # [TODO] figure out how to deal with mimic joints in pinocchio + # breakpoint() + + # return shared variables + return ( + mj_model, + mjx_model, + mjx_data, + mj_joint_names, + pin_model, + pin_data, + pin_joint_names, + qpos_mjx, + qpos_pin, + mj2pin, + pin2mj, + ) + + +def test_joint_orders(shared_variables): + """Tests joint ordering.""" + mj_model = shared_variables[0] + mj_joint_names = shared_variables[3] + pin_model = shared_variables[4] + pin_joint_names = shared_variables[6] + mj2pin, pin2mj = get_joint_orders(mj_model, pin_model) + 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 + + +def test_fk(shared_variables): + """Tests forward kinematics. Also implicitly tests conversion functions between mjx and pin.""" + # loading global variables + ( + mj_model, + mjx_model, + mjx_data, + mj_joint_names, + pin_model, + pin_data, + pin_joint_names, + qpos_mjx, + qpos_pin, + mj2pin, + pin2mj, + ) = shared_variables + + # running FK test + mjx_data = mjx.kinematics(mjx_model, mjx_data.replace(qpos=qpos_mjx)) + pin.forwardKinematics(pin_model, pin_data, qpos_pin) + + # testing in one direction (mjx coords -> pin) + xanchors_pin = mjx_xanchor_to_pin(mj2pin, mjx_data) + for i, x in enumerate(xanchors_pin): + assert np.allclose( + np.array(x), + np.array(pin_data.oMi[i].translation), + rtol=1e-6, + atol=1e-6, # loosen atol a bit + ) + + # testing other direction (pin coords -> mjx) + xanchors_mjx = mjx_data.xanchor + xanchors_mjx_converted = pin_xanchor_to_mjx(pin2mj, pin_data) + for i, xanchor_mjx in enumerate(xanchors_mjx): + if mj2pin[i] is not None: + assert np.allclose( + np.array(xanchor_mjx), + np.array(xanchors_mjx_converted[i, :]), + rtol=1e-6, + atol=1e-6, # loosen atol a bit + ) + + +# def test_fd(shared_variables): +# """Tests forward dynamics.""" +# mj_model, mjx_model, mjx_data, mj_joint_names, pin_model, pin_data, pin_joint_names, qpos_mjx, qpos_pin, mj2pin, pin2mj = shared_variables + +# # running FD test +# dv_mjx = np.array(mjx_data.qacc) +# mjx_data = mjx.forward(mjx_model, mjx_data.replace(qpos=qpos_mjx)) + +# M = pin.crba(pin_model, pin_data, qpos_pin) +# qvel_pin = np.zeros(pin_model.nv) # [TODO] make this more exotic +# zero_accel = np.zeros(pin_model.nv) +# tau0 = pin.rnea(pin_model, pin_data, qpos_pin, qvel_pin, zero_accel) +# # [TODO] add applied forces tau to pinocchio and mjx for the test, check about gravity diff --git a/tests/test_model_io.py b/tests/test_model_io.py index 5288ba8f..b2773610 100644 --- a/tests/test_model_io.py +++ b/tests/test_model_io.py @@ -1,13 +1,12 @@ from pathlib import Path import igl -import mujoco as mj import numpy as np import trimesh from dm_control import mjcf from lxml import etree -from mujoco import mjx +import mujoco as mj from ambersim import ROOT from ambersim.utils._internal_utils import _rmtree from ambersim.utils.conversion_utils import convex_decomposition_file, save_model_xml @@ -18,6 +17,7 @@ load_mjx_model_and_data_from_file, mj_to_mjx_model_and_data, ) +from mujoco import mjx def test_load_model():