Visualization

MJINX provides utilities for visualizing robots, problem components, and optimization results. These visualization tools help with debugging, analysis, and presentation of inverse kinematics solutions.

The visualization module offers: - Functions for rendering robot states - Tools for visualizing task errors and barrier values - Methods for displaying optimization trajectories - Integration with popular plotting libraries

Using visualization alongside MJINX’s computational capabilities allows for both development and deployment of robust inverse kinematics solutions.

class MarkerData(name, id, type, size, pos=<factory>, rot=<factory>, rgba=<factory>)

A dataclass for storing marker data.

Parameters:

  • name (str) – The name of the marker.

  • id (int) – The unique identifier of the marker.

  • type (mjtGeom) – The type of geometry for the marker.

  • size (ndarray) – The size of the marker.

  • pos (ndarray) – The position of the marker in 3D space. Defaults to [0, 0, 0].

  • rot (ndarray) – The rotation of the marker. Defaults to [1, 0, 0, 0] (identity quaternion).

  • rgba (ndarray) – The RGBA color values of the marker. Defaults to [0.5, 0.5, 0.5, 0.3].

property rot_matrix: ndarray

Returns a raveled rotation matrix, generated from orientation.

It checks the ndim of the rot field. If ndim=1, it assumes quaternion (scalar first) is passed, if ndim=2, it assumes that rotation matrix is passed.

Returns:

Raveled rotation matrix with shape (9,).

Raises:

ValueError – If the rotation data has invalid dimensions or length.

class BatchVisualizer(model_path, n_models, geom_group=2, alpha=0.5, record=False, passive_viewer=True, filename='', record_res=(1024, 1024))

A class for batch visualization of multiple MuJoCo model instances.

This class allows for the visualization of multiple instances of a given model, with customizable transparency and marker options. It also supports recording the visualization as a video.

Parameters:

  • model_path (str) – Path to the MuJoCo model file.

  • n_models (int) – Number of model instances to visualize.

  • geom_group (int) – Geometry group to render, defaults to 2.

  • alpha (float) – Transparency value for the models, defaults to 0.5.

  • record (bool) – If True, records and saves mp4 scene recording, defaults to False.

  • filename (str) – Name of the file to save without extension, defaults to current datetime.

  • record_res (tuple[int, int]) – Resolution of recorded video (width, height), defaults to (1024, 1024).

remove_high_level_body_tags(mjcf_str, model_directory)

Remove high-level body tags from the MJCF XML string and process mesh assets.

This method modifies the XML structure by moving children of high-level body tags directly under the worldbody tag. It also processes mesh elements, updating their file attributes and collecting asset data.

Note that this solution is obviously very hacky and dirty, however the author considers this easier and (barely) enough to complete the required task. The attempts to solve this using dm_control failed, see: https://github.com/google-deepmind/dm_control/issues/407

Parameters:

  • mjcf_str (str) – The MJCF XML string to process.

  • model_directory (str) – The directory containing the model and its assets.

Return type:

tuple[str, dict[str, bytes]]

Returns:

A tuple containing the modified XML string and a dictionary of asset data.

add_markers(name, size, marker_alpha, color_begin, color_end=None, marker_type=<mjtGeom.mjGEOM_SPHERE: 2>, n_markers=1)

Add markers to the visualization.

Parameters:

  • name (str | Sequence[str]) – Name or sequence of names for the markers.

  • size (ndarray | Array | float) – Size of the markers. Can be a single float or an array of size 3.

  • marker_alpha (float) – Transparency of the markers.

  • color_begin (ndarray) – Starting color for marker interpolation.

  • color_end (ndarray | None) – Ending color for marker interpolation. If None, uses color_begin.

  • marker_type (mjtGeom) – Type of marker geometry, defaults to sphere.

  • n_markers (int) – Amount of markers to add. Defaults to 1.

get_prefix(i)

Generate a prefix for the i-th model instance.

Parameters:

i (int) – Index of the model instance.

Return type:

str

Returns:

Prefix string for the model.

update(q)

Update the model positions and record frame if enabled.

Parameters:

q (ndarray | Array) – Array of joint positions for all model instances.

save_video(fps)

Save the recorded frames as an MP4 video.

Parameters:

fps (float) – Frames per second for the output video.

close()

Close the viewer and clean up resources.