Improve support of visualizing heightmaps/hfield in the Mujoco passive viewer

[diegoferigo]

I couldn't make the previous logic work correctly. This PR makes sure of the following:

• The converted MjModel is initialized with a hfield grid whose number of rows and columns can be customized.
• The area over which the heightmap is sampled can now be defined.
• The MJCF model is dynamically updated with the wideness and elevation of the sampled heightmap data.

Here below the result.

Note that if the heightmap would be rendered incorrectly, the sphere would not lay over the surface as it would mean that the simulated data and the visualized terrain differ. The image shows that the two match as expected.

MWE

import time

import jax.numpy as jnp
import jax_dataclasses
import jaxsim
import jaxsim.api as js
import numpy as np
import rod.builder.primitives
from jaxsim import typing as jtp


# Define a custom terrain:
# https://it.mathworks.com/help/matlab/ref/peaks.html
@jax_dataclasses.pytree_dataclass
class PeaksTerrain(jaxsim.terrain.Terrain):

    def height(self, x: jtp.FloatLike, y: jtp.FloatLike) -> jtp.Float:

        z = (
            3 * (1 - x) ** 2 * jnp.exp(-(x**2) - (y + 1) ** 2)
            - 10 * (x / 5 - x**3 - y**5) * jnp.exp(-(x**2) - y**2)
            - 1 / 3 * jnp.exp(-((x + 1) ** 2) - y**2)
        )

        return jnp.array(0.25 * z, dtype=float)


# Build the ROD model of a sphere.
sphere = (
    rod.builder.primitives.SphereBuilder(name="sphere", mass=1.0, radius=0.1)
    .build_model()
    .add_link(name="link")
    .add_inertial()
    .add_visual()
    .add_collision()
    .build()
)


# Build the JaxSim model.
model = js.model.JaxSimModel.build_from_model_description(
    model_description=sphere,
    terrain=PeaksTerrain(),
)

# Initialize the sphere with the origin be located over the terrain at a
# given (x, y) coordinates.
x, y = -0.42, -0.62
data = js.data.JaxSimModelData.build(
    model=model,
    base_position=jnp.array([x, y, model.terrain.height(x=x, y=y)]),
)


# ================
# 3D visualization
# ================

import jaxsim.mujoco

mjcf_string, assets = jaxsim.mujoco.RodModelToMjcf().convert(
    rod_model=model.built_from,
    heightmap=True,
    heightmap_samples_xy=(256, 256),
)

mj_model_helper = jaxsim.mujoco.MujocoModelHelper.build_from_xml(
    mjcf_description=mjcf_string,
    assets=assets,
    heightmap=lambda x, y: model.terrain.height(x=x, y=y).tolist(),
    heightmap_radius_xy=(5.0, 5.0),
)

viz = jaxsim.mujoco.MujocoVisualizer(
    model=mj_model_helper.model, data=mj_model_helper.data
)


with viz.open() as viewer:

    with viewer.lock():
        mj_model_helper.set_base_position(position=np.array(data.base_position()))
        mj_model_helper.set_base_orientation(
            orientation=np.array(data.base_orientation(dcm=False))
        )

    viz.sync(viewer=viewer)

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

---

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

[flferretti]

LGTM, thanks!