Add support for inclined planes and heightmap visualization

[flferretti]

This pull request adds support for inclined planes and heightmap visualization using the MuJoCo visualizer.

Default behavior:

video.mp4

To run the examples:

MWE Inclined Plane

import pathlib

import jax.numpy as jnp
import numpy as np
import resolve_robotics_uri_py
import rod

import jaxsim
import jaxsim.api as js
from jaxsim import VelRepr, integrators

# Find the urdf file.
urdf_path = resolve_robotics_uri_py.resolve_robotics_uri(uri="file://stickbot.urdf")
# Build the ROD model.
rod_sdf = rod.Sdf.load(sdf=urdf_path)

# Create a terrain.
plane_normal = jnp.array([0.0, -0.1, 0.2])
terrain = jaxsim.terrain.PlaneTerrain(
    plane_normal=plane_normal,
)

# Build the model.
model = js.model.JaxSimModel.build_from_model_description(
    model_description=rod_sdf.model,
    terrain=terrain,
)


# Build the model's data.
# Set already here the initial base position.
data0 = js.data.JaxSimModelData.build(
    model=model,
    base_position=jnp.array([0, 0, 0.85]),
    velocity_representation=VelRepr.Inertial,
)

# Update the soft-contact parameters.
# By default, only 1 support point is used as worst-case scenario.
# Feel free to tune this with more points to get a less stiff system.
data0 = data0.replace(
    soft_contacts_params=js.contact.estimate_good_soft_contacts_parameters(
        model, number_of_active_collidable_points_steady_state=2
    )
)

# =====================
# Create the integrator
# =====================

# Create a RK4 integrator integrating the quaternion on SO(3).
integrator = integrators.fixed_step.RungeKutta4SO3.build(
    dynamics=js.ode.wrap_system_dynamics_for_integration(
        model=model,
        data=data0,
        system_dynamics=js.ode.system_dynamics,
    ),
)

# =========================================
# Visualization in Mujoco viewer / renderer
# =========================================

from jaxsim.mujoco import MujocoModelHelper, MujocoVideoRecorder, RodModelToMjcf

# Convert the ROD model to a Mujoco model.
mjcf_string, assets = RodModelToMjcf.convert(
    rod_model=rod_sdf.models()[0],
    considered_joints=list(model.joint_names()),
    plane_normal=plane_normal,
)

# Build the Mujoco model helper.
mj_model_helper = self = MujocoModelHelper.build_from_xml(
    mjcf_description=mjcf_string, assets=assets
)

# Create the video recorder.
recorder = MujocoVideoRecorder(
    model=mj_model_helper.model,
    data=mj_model_helper.data,
    fps=int(1 / 0.010),
    width=320 * 4,
    height=240 * 4,
)

# ==============
# Recording loop
# ==============

# Initialize the integrator.
t0 = 0.0
tf = 5.0
dt = 0.001_000
integrator_state = integrator.init(x0=data0.state, t0=t0, dt=dt)

# Initialize the loop.
data = data0.copy()
joint_names = list(model.joint_names())

while data.time_ns < tf * 1e9:

    # Integrate the dynamics.
    data, integrator_state = js.model.step(
        dt=dt,
        model=model,
        data=data,
        integrator=integrator,
        integrator_state=integrator_state,
        # Optional inputs
        joint_forces=None,
        external_forces=None,
    )

    # Extract the generalized position.
    s = data.state.physics_model.joint_positions
    W_p_B = data.state.physics_model.base_position
    W_Q_B = data.state.physics_model.base_quaternion

    # Update the data object stored in the helper, which is shared with the recorder.
    mj_model_helper.set_base_position(position=np.array(W_p_B))
    mj_model_helper.set_base_orientation(orientation=np.array(W_Q_B), dcm=False)
    mj_model_helper.set_joint_positions(positions=np.array(s), joint_names=joint_names)

    # Record the frame if the time is right to get the desired fps.
    if data.time_ns % jnp.array(1e9 / recorder.fps).astype(jnp.uint64) == 0:
        recorder.record_frame(camera_name=None)

# Store the video.
video_path = pathlib.Path("~/video.mp4").expanduser()
recorder.write_video(path=video_path, exist_ok=True)

# Clean up the recorder.
recorder.frames = []
recorder.renderer.close()

video.mp4

MWE Heightmap

# Coming soon

C.C. @lorycontixd

---

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

[lorycontixd]

@flferretti Che versione di Jaxsim stai usando?
Nel pull del branch functional non mi risulta che la funzione JaxSimModel.build_from_model_description prenda in ingresso un terrain.

[flferretti]

@lorycontixd you can run gh pr checkout 122 to checkout the branch relative to this PR

[diegoferigo]

@flferretti this PR has been automatically closed since its target branch was merged. I always forget this behavior. Feel free to open it again against main.

[flferretti]

Yes, it is a pretty annoying behavior. Reopening in #125