visualizers.mujoco.model#

Module Contents#

visualizers.mujoco.model.DEFAULT_GRID_SECTION_COLOR = (0.02, 0.14, 0.44)#
visualizers.mujoco.model.DEFAULT_GRID_CELL_COLOR = (0.27, 0.55, 1)#
visualizers.mujoco.model.DEFAULT_SPHERE_SUBDIVISIONS = 3#
visualizers.mujoco.model.DEFAULT_SPLINE_COLOR = (0.8, 0.1, 0.8)#
visualizers.mujoco.model.DEFAULT_BEST_SPLINE_COLOR = (0.96, 0.7, 0.0)#
class visualizers.mujoco.model.ViserMjModel(target: viser.ViserServer | viser.ClientHandle, model: mujoco.MjModel, show_ground_plane: bool = True, alpha_scale: float | None = None, namespace: str | None = None)#

Helper for rendering MJCF models in viser.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) – ViserServer or ClientHandle to add MjModel to.

  • model (mujoco.MjModel) – MjModel to be visualized.

  • data – MjData of model in initial configuration.

  • show_ground_plane (bool) –

  • alpha_scale (float | None) –

  • namespace (str | None) –

add_geom(geom_name: str, geom: Any) → None#

Helper function for adding geoms to scene tree.

Parameters:

  • geom_name (str) –

  • geom (Any) –

Return type:

None

get_color(geom: Any) → numpy.ndarray#
Parameters:

geom (Any) –

Return type:

numpy.ndarray

add_traces(num_traces: int = 0, all_traces_rollout_size: int = 0, trace_name: str = 'trace') → None#

Add a collection of all traces to the visualizer, done in one go to avoid having too many things.

We have two sets of traces to care about: the “elite” reward traces and the regular ones. Due to how the line segments work, we only need one handle per type.

Parameters:

  • num_traces (int) –

  • all_traces_rollout_size (int) –

  • trace_name (str) –

Return type:

None

remove_traces() → None#

Remove traces.

Return type:

None

set_data(data: mujoco.MjData) → None#

Write updated configuration from mujoco data to viser viewer.

Parameters:

data (mujoco.MjData) –

Return type:

None

set_traces(traces: numpy.ndarray | None, all_traces_rollout_size: int) → None#

Write updated traces to viser viewer.

Parameters:

  • traces (numpy.ndarray | None) – trace sensors readings of size (self.num_elite * all_traces_rollout_size, 2, 3)

  • all_traces_rollout_size (int) – num_trace_sensors * single_rollout, size of all grouped trace sensor rollouts

Return type:

None

remove() → None#

Wrapper function to remove all geometries from Viser.

Return type:

None

visualizers.mujoco.model.add_plane(target: viser.ViserServer | viser.ClientHandle, name: str, pos: Tuple[float, float, float] | numpy.ndarray = (0.0, 0.0, 0.0), quat: Tuple[float, float, float, float] | numpy.ndarray = (1.0, 0.0, 0.0, 0.0)) → viser.SceneNodeHandle#

Add a plane geometry to the visualizer with optional position, quaternion, material, and name.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • pos (Tuple[float, float, float] | numpy.ndarray) –

  • quat (Tuple[float, float, float, float] | numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_sphere(target: viser.ViserServer | viser.ClientHandle, name: str, radius: float, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add a sphere geometry to the visualizer with optional position, quaternion, material, and name.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • radius (float) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_cylinder(target: viser.ViserServer | viser.ClientHandle, name: str, radius: float, height: float, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add a cylinder geometry to the visualizer with optional position, quaternion, material, and name.

The cylinder is aligned with the z-axis

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • radius (float) –

  • height (float) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_box(target: viser.ViserServer | viser.ClientHandle, name: str, size: numpy.ndarray, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add a box geometry to the visualizer with optional position, quaternion, material, and name.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • size (numpy.ndarray) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_capsule(target: viser.ViserServer | viser.ClientHandle, name: str, radius: float, length: float, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add a capsule geometry to the visualizer with optional position, quaternion, material, and name.

The capsule is aligned with the z-axis

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • radius (float) –

  • length (float) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_ellipsoid(target: viser.ViserServer | viser.ClientHandle, name: str, scaling: numpy.ndarray, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add an ellipsoid geometry to the visualizer.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • scaling (numpy.ndarray) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_mesh(target: viser.ViserServer | viser.ClientHandle, name: str, vertices: numpy.ndarray, faces: numpy.ndarray, pos: numpy.ndarray, quat: numpy.ndarray, rgba: numpy.ndarray) → viser.SceneNodeHandle#

Add a triangular mesh geometry to the visualizer

Add a triangular mesh geometry to the visualizer with specified vertices and faces, with optional position, quaternion, material, and name.

Vertices: float (N, 3) and faces: int (M, 3).

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • vertices (numpy.ndarray) –

  • faces (numpy.ndarray) –

  • pos (numpy.ndarray) –

  • quat (numpy.ndarray) –

  • rgba (numpy.ndarray) –

Return type:

viser.SceneNodeHandle

visualizers.mujoco.model.add_spline(target: viser.ViserServer | viser.ClientHandle, name: str, positions: tuple[tuple[float, float, float], Ellipsis] | numpy.ndarray, pos: Tuple[float, float, float] | numpy.ndarray = (0.0, 0.0, 0.0), quat: Tuple[float, float, float, float] | numpy.ndarray = (1.0, 0.0, 0.0, 0.0), rgb: Tuple[float, float, float] = DEFAULT_SPLINE_COLOR, line_width: float = 4.0, segments: int | None = None, visible: bool = True) → viser.SplineCatmullRomHandle#

Add a spline to the visualizer with optional position, quaternion.

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) –

  • name (str) –

  • positions (tuple[tuple[float, float, float], Ellipsis] | numpy.ndarray) –

  • pos (Tuple[float, float, float] | numpy.ndarray) –

  • quat (Tuple[float, float, float, float] | numpy.ndarray) –

  • rgb (Tuple[float, float, float]) –

  • line_width (float) –

  • segments (int | None) –

  • visible (bool) –

Return type:

viser.SplineCatmullRomHandle

visualizers.mujoco.model.add_segments(target: viser.ViserServer | viser.ClientHandle, name: str, points: numpy.ndarray, pos: Tuple[float, float, float] | numpy.ndarray = (0.0, 0.0, 0.0), quat: Tuple[float, float, float, float] | numpy.ndarray = (1.0, 0.0, 0.0, 0.0), rgb: Tuple[float, float, float] = DEFAULT_SPLINE_COLOR, line_width: float = 4.0, visible: bool = True) → viser.LineSegmentsHandle#

Add line segments to the visualizer with an optional position and orientation.

TODO(@bhung) Potentially add support for different kinds of segments

Parameters:

  • target (viser.ViserServer | viser.ClientHandle) – ViserServer or handle to attach the segments to

  • name (str) – name of the segments

  • points (numpy.ndarray) – size (N x 2 x 3) where index 0 is point, 1 is start vs end, and 2 is 3D coord

  • pos (Tuple[float, float, float] | numpy.ndarray) – position that the points are defined with respect to. Defaults to origin

  • quat (Tuple[float, float, float, float] | numpy.ndarray) – orientation that the points are defined with respect to. Defaults to identity

  • rgb (Tuple[float, float, float]) – colors of the points. Can be sized (N x 2 x 3) or a broadcastable shape

  • line_width (float) – width of the line, in pixels

  • visible (bool) – whether or not the lines are initially visible

Return type:

viser.LineSegmentsHandle

visualizers.mujoco.model.rgba_float_to_int(rgba_float: numpy.ndarray) → numpy.ndarray#
Parameters:

rgba_float (numpy.ndarray) –

Return type:

numpy.ndarray

visualizers.mujoco.model.rgba_int_to_float(rgba_int: numpy.ndarray) → numpy.ndarray#
Parameters:

rgba_int (numpy.ndarray) –

Return type:

numpy.ndarray

visualizers.mujoco.model.set_mesh_color(mesh: trimesh.Trimesh, rgba: numpy.ndarray) → None#
Parameters:

  • mesh (trimesh.Trimesh) –

  • rgba (numpy.ndarray) –

Return type:

None

visualizers.mujoco.model.set_spline_positions(handle: viser.SplineCatmullRomHandle, positions: tuple[tuple[float, float, float], Ellipsis] | numpy.ndarray) → None#

Set the spline waypoints.

Parameters:

  • handle (viser.SplineCatmullRomHandle) –

  • positions (tuple[tuple[float, float, float], Ellipsis] | numpy.ndarray) –

Return type:

None

visualizers.mujoco.model.set_segment_points(handle: viser.LineSegmentsHandle, points: numpy.ndarray) → None#

Set the line waypoints.

Parameters:

  • handle (viser.LineSegmentsHandle) – handle for the line segments

  • points (numpy.ndarray) – start and end points of the line segments (N x 2 x 3)

Return type:

None