Add helper to convert JaxSimModelData to mujoco.mjData

[diegoferigo]

The main benefit of using this helper w.r.t. plain usage of MujocoModelHelper is the logic to initialize the positions of removed joints. An example usage is the following:

import time

import jax.numpy as jnp
import jaxsim.api as js
import jaxsim.mujoco
import numpy as np
import requests

repo = "icub-tech-iit/ergocub-gazebo-simulations"
url = f"https://raw.githubusercontent.com/{repo}/master/models/stickBot/model.urdf"

# Download and read file.
urdf_string = requests.get(url).text

# Build the model.
model = js.model.JaxSimModel.build_from_model_description(model_description=urdf_string)

# Reduce the model by specifying custom elbow angles.
model = js.model.reduce(
    model=model,
    considered_joints=tuple(
        j
        for j in model.joint_names()
        # Remove upper body and sensors.
        if "neck" not in j
        and "wrist" not in j
        and "torso" not in j
        and "elbow" not in j
        and "shoulder" not in j
        and "lidar" not in j
        and "camera" not in j
    ),
    locked_joint_positions={
        "r_elbow": np.deg2rad(45),
        "l_elbow": np.deg2rad(45),
    },
)

# Create the model data.
data = js.data.JaxSimModelData.build(
    model=model,
    base_position=jnp.array([0, 0, 2.0]),
)

# ===================
# Open the visualizer
# ===================

# Convert the model to MJCF.
mjcf_string, assets = jaxsim.mujoco.UrdfToMjcf().convert(urdf=model.built_from)

# Create the Mujoco model helper.
mj_model_helper = jaxsim.mujoco.MujocoModelHelper.build_from_xml(
    mjcf_description=mjcf_string,
    assets=assets,
)

# Store in the model helper the data initialized from the JaxSim model.
# This is helpful the first time, so that the removed joint positions are
# stored in the Mujoco data object.
mj_model_helper.data = jaxsim.mujoco.mujoco_data_from_jaxsim(
    mujoco_model=mj_model_helper.model,
    jaxsim_model=model,
    jaxsim_data=data,
)

# Create the visualizer.
viz = jaxsim.mujoco.MujocoVisualizer(
    model=mj_model_helper.model, data=mj_model_helper.data
)


with viz.open(
    lookat=np.array(data.base_position()), elevation=-10, distance=5, azimut=180 - 30
) as viewer:

    # Here you can loop over a simulated trajectory.
    # It is not necessary to use the jaxsim.mujoco.mujoco_data_from_jaxsim helper.

    # for ...:
    #     with viewer.lock():
    #         mj_model_helper.set_base_position(position=np.array(...))
    #         mj_model_helper.set_base_orientation(orientation=np.array(...))
    #
    #         if model.dofs() > 0:
    #             mj_model_helper.set_joint_positions(
    #                 joint_names=list(model.joint_names()), positions=np.array(...)
    #             )

    viz.sync(viewer=viewer)

    while viewer.is_running():
        time.sleep(0.500)

---

📚 Documentation preview 📚: https://jaxsim--264.org.readthedocs.build//264/

[diegoferigo]

cc @xela-95

[flferretti]

Thanks a lot Diego! I left a couple comments

[xela-95]

This is cool thanks! 🚀

[flferretti]

Thanks Diego! Good to go after the last suggestion