diff --git a/src/mjviser/conversions.py b/src/mjviser/conversions.py index ad5b284..18a052b 100644 --- a/src/mjviser/conversions.py +++ b/src/mjviser/conversions.py @@ -44,8 +44,15 @@ def _is_cubemap_texture(mj_model: mujoco.MjModel, texid: int) -> bool: return h == w * 6 and nc in (1, 3, 4) -# Cube map face order with the outward scene axis each face sits on: -# +X (right), -X (left), +Y (up), -Y (down), +Z (front), -Z (back). +def _is_2d_texture_supported(mj_model: mujoco.MjModel, texid: int) -> bool: + """Return True if texid is a 2D texture with a channel count we can extract.""" + return int(mj_model.tex_nchannel[texid]) in (1, 3, 4) + + +# MuJoCo stores cube faces in the order right, left, up, down, front, back and +# uploads them to GL_TEXTURE_CUBE_MAP_POSITIVE_X + i. For a geom (regular) cube +# texture MuJoCo samples with texcoords = the geom-local position (x, y, z), so +# in the geom frame the faces sit on +X, -X, +Y, -Y, +Z, -Z respectively. _CUBEMAP_AXES: np.ndarray = np.array( [[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]], dtype=np.float64, @@ -292,11 +299,11 @@ def create_primitive_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Tri if geom_type == mjtGeom.mjGEOM_HFIELD: return _create_heightfield_mesh(mj_model, geom_id) - # Box primitives with a cube map texture get per-face textured rendering. - if geom_type == mjtGeom.mjGEOM_BOX: - textured = _create_cubemap_box_mesh(mj_model, geom_id) - if textured is not None: - return textured + # Textured primitives (box/sphere/ellipsoid cube maps, plane 2D textures) are + # dispatched in one place so this and get_geom_texture_id stay in sync. + textured = _textured_primitive_mesh(mj_model, geom_id) + if textured is not None: + return textured if geom_type == mjtGeom.mjGEOM_PLANE: size = mj_model.geom_size[geom_id] @@ -310,19 +317,48 @@ def create_primitive_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Tri return mesh -# Per cube face, the (image-right, image-up) directions in scene coordinates as -# seen by a viewer looking at the face from outside the cube. Indexed by the -# face order in _CUBEMAP_AXES. Matches MuJoCo's cube map convention. -_CUBEMAP_FACE_BASIS: tuple[tuple[np.ndarray, np.ndarray], ...] = ( - (np.array([0.0, -1.0, 0.0]), np.array([0.0, 0.0, 1.0])), # +X - (np.array([0.0, 1.0, 0.0]), np.array([0.0, 0.0, 1.0])), # -X - (np.array([1.0, 0.0, 0.0]), np.array([0.0, 0.0, -1.0])), # +Y - (np.array([1.0, 0.0, 0.0]), np.array([0.0, 0.0, 1.0])), # -Y - (np.array([1.0, 0.0, 0.0]), np.array([0.0, 1.0, 0.0])), # +Z - (np.array([-1.0, 0.0, 0.0]), np.array([0.0, 1.0, 0.0])), # -Z +# OpenGL cube-map per-face selectors, indexed by face order +# (right, left, up, down, front, back). For a point/direction p on face fi: +# sc = sc_sign * p[sc_axis]; tc = tc_sign * p[tc_axis]; ma = |p[major]| +# s = (sc/ma + 1)/2; t = (tc/ma + 1)/2 (t = 0 is the top row of the face) +# This is the convention MuJoCo's renderer uses for a geom cube texture; the box +# and sphere both route through it so they stay consistent. +_CUBE_FACE_ST: tuple[tuple[int, int, int, int, int], ...] = ( + (2, -1, 1, -1, 0), # +X right: sc=-z, tc=-y + (2, 1, 1, -1, 0), # -X left: sc=+z, tc=-y + (0, 1, 2, 1, 1), # +Y up: sc=+x, tc=+z + (0, 1, 2, -1, 1), # -Y down: sc=+x, tc=-z + (0, 1, 1, -1, 2), # +Z front: sc=+x, tc=-y + (0, -1, 1, -1, 2), # -Z back: sc=-x, tc=-y ) +def _cube_face_st( + positions: np.ndarray, faces: np.ndarray +) -> tuple[np.ndarray, np.ndarray]: + """Return per-point (s, t) in [0,1] for points lying on the given cube faces.""" + table = np.array(_CUBE_FACE_ST) + sc_axis, sc_sign, tc_axis, tc_sign, major = (table[faces, k] for k in range(5)) + idx = np.arange(len(faces)) + sc = sc_sign * positions[idx, sc_axis] + tc = tc_sign * positions[idx, tc_axis] + ma = np.maximum(np.abs(positions[idx, major]), 1e-12) + return (sc / ma + 1.0) * 0.5, (tc / ma + 1.0) * 0.5 + + +def _gl_cube_face(directions: np.ndarray) -> np.ndarray: + """Return the cube face index (right, left, up, down, front, back) per direction.""" + x, y, z = directions[:, 0], directions[:, 1], directions[:, 2] + ax, ay, az = np.abs(x), np.abs(y), np.abs(z) + x_major = (ax >= ay) & (ax >= az) + y_major = (ay > ax) & (ay >= az) + return np.where( + x_major, + np.where(x >= 0, 0, 1), + np.where(y_major, np.where(y >= 0, 2, 3), np.where(z >= 0, 4, 5)), + ).astype(np.int64) + + def _extract_cubemap_atlas(mj_model: mujoco.MjModel, texid: int) -> Image.Image | None: """Build a vertical-strip PIL atlas from a cube map texture. @@ -341,9 +377,10 @@ def _extract_cubemap_atlas(mj_model: mujoco.MjModel, texid: int) -> Image.Image atlas = Image.new(mode, (w, 6 * w)) for i in range(6): - # MuJoCo stores rows bottom-up; flip to PIL top-down so the face - # image appears right-side-up in the atlas. - arr = np.flipud(data[i]).astype(np.uint8) + # Face rows are stored top-down (PNG order); face i goes to atlas rows + # [i*w, (i+1)*w). The box UVs sample this with the same (s, t) the sphere + # sampler uses, so both share the OpenGL cube-map orientation. + arr = data[i].astype(np.uint8) if nc == 1: arr = arr.reshape(w, w) atlas.paste(Image.fromarray(arr, mode=mode), (0, i * w)) @@ -370,36 +407,32 @@ def _create_cubemap_box_mesh( if atlas is None: return None - sx, sy, sz = mj_model.geom_size[geom_id] - centers = _CUBEMAP_AXES * np.array([sx, sx, sy, sy, sz, sz])[:, None] - half_extents = [(sy, sz), (sy, sz), (sx, sz), (sx, sz), (sx, sy), (sx, sy)] + size = mj_model.geom_size[geom_id] + combos = ((-1, -1), (1, -1), (1, 1), (-1, 1)) verts = np.zeros((24, 3), dtype=np.float64) uvs = np.zeros((24, 2), dtype=np.float64) faces = np.zeros((12, 3), dtype=np.int64) - for fi in range(6): - c = centers[fi] - r, u = _CUBEMAP_FACE_BASIS[fi] - hr, hu = half_extents[fi] - # Corners in image-relative order: bottom-left, bottom-right, top-right, - # top-left. + for fi, (_sc_axis, _sc_sign, _tc_axis, _tc_sign, major) in enumerate(_CUBE_FACE_ST): + n = _CUBEMAP_AXES[fi] + free = [k for k in range(3) if k != major] base = fi * 4 - verts[base + 0] = c - hr * r - hu * u - verts[base + 1] = c + hr * r - hu * u - verts[base + 2] = c + hr * r + hu * u - verts[base + 3] = c - hr * r + hu * u - # Atlas UVs (OpenGL convention: v=0 at bottom of atlas image). - v_top = 1.0 - fi / 6.0 - v_bot = 1.0 - (fi + 1) / 6.0 - uvs[base + 0] = (0.0, v_bot) - uvs[base + 1] = (1.0, v_bot) - uvs[base + 2] = (1.0, v_top) - uvs[base + 3] = (0.0, v_top) - # Pick triangle winding so the face normal points outward. The image - # basis (r, u) is chosen so the texture reads correctly when viewed - # from outside, which means r × u may point either way relative to - # the outward axis depending on the face. - if np.dot(np.cross(r, u), _CUBEMAP_AXES[fi]) >= 0: + corners = np.zeros((4, 3)) + for ci, (a, b) in enumerate(combos): + corners[ci, major] = n[major] * size[major] + corners[ci, free[0]] = a * size[free[0]] + corners[ci, free[1]] = b * size[free[1]] + verts[base : base + 4] = corners + + # Atlas UVs from the shared GL (s, t): face fi spans atlas rows + # [fi*w, (fi+1)*w), so v = 1 - (fi + t) / 6 (glTF v-up over the strip). + s, t = _cube_face_st(corners, np.full(4, fi)) + uvs[base : base + 4, 0] = s + uvs[base : base + 4, 1] = 1.0 - (fi + t) / 6.0 + + # Wind triangles so the face normal points outward. + normal = np.cross(corners[1] - corners[0], corners[2] - corners[0]) + if np.dot(normal, n) >= 0: faces[fi * 2 + 0] = (base + 0, base + 1, base + 2) faces[fi * 2 + 1] = (base + 0, base + 2, base + 3) else: @@ -418,6 +451,208 @@ def _create_cubemap_box_mesh( return mesh +def _cubemap_sample_colors( + mj_model: mujoco.MjModel, texid: int, directions: np.ndarray +) -> np.ndarray | None: + """Sample a cube map along directions, returning uint8 RGBA. + + Implements the OpenGL cube-map lookup MuJoCo uses for a geom (regular) + cube texture: texcoords are the geom-local direction (x, y, z), faces + are stored right, left, up, down, front, back, and the per-face (s, t) + follow the GL spec. Face rows are stored top-down (PNG order), so t=0 + reads the top of the face image. + """ + if not _is_cubemap_texture(mj_model, texid): + return None + w = int(mj_model.tex_width[texid]) + nc = int(mj_model.tex_nchannel[texid]) + adr = int(mj_model.tex_adr[texid]) + data = mj_model.tex_data[adr : adr + 6 * w * w * nc].reshape(6, w, w, nc) + + face = _gl_cube_face(directions) + s, t = _cube_face_st(directions, face) + col = np.clip((s * w).astype(int), 0, w - 1) + row = np.clip((t * w).astype(int), 0, w - 1) + px = data[face, row, col] + + out = np.full((len(directions), 4), 255, dtype=np.uint8) + if nc == 1: + out[:, :3] = px[:, :1] + else: + out[:, :nc] = px[:, :nc] + return out + + +def _cubemap_to_equirect( + mj_model: mujoco.MjModel, + texid: int, + scale: np.ndarray, + width: int = 512, + height: int = 256, +) -> Image.Image | None: + """Bake a cube map into an equirectangular (lat-long) RGBA image. + + Row 0 is the north pole (+Z); columns sweep longitude 0..2*pi with 0 along + +X. Sampling reuses _cubemap_sample_colors so the face layout matches the box + renderer. The per-axis ``scale`` matches the geom's size: MuJoCo samples a + geom cube texture by the geom-local position, so an ellipsoid's faces land at + different angles than a sphere's. + """ + lon = (np.arange(width) + 0.5) / width * (2.0 * np.pi) + lat = (np.arange(height) + 0.5) / height * np.pi + lon_g, lat_g = np.meshgrid(lon, lat) + sin_lat = np.sin(lat_g) + directions = np.stack( + [sin_lat * np.cos(lon_g), sin_lat * np.sin(lon_g), np.cos(lat_g)], axis=-1 + ).reshape(-1, 3) + colors = _cubemap_sample_colors(mj_model, texid, directions * scale) + if colors is None: + return None + return Image.fromarray(colors.reshape(height, width, 4), "RGBA") + + +def _create_textured_sphere_mesh( + mj_model: mujoco.MjModel, geom_id: int +) -> trimesh.Trimesh | None: + """Build a UV-textured sphere or ellipsoid mesh from a cube map. + + Returns None when the geom has no material or no cube map texture. The + cube map is baked into an equirectangular image and applied to a + lat-long sphere with a duplicated seam column, so the texture stays + crisp regardless of mesh density. + """ + matid = int(mj_model.geom_matid[geom_id]) + if matid < 0 or matid >= mj_model.nmat: + return None + texid = _get_texture_id(mj_model, matid) + if texid < 0 or not _is_cubemap_texture(mj_model, texid): + return None + size = mj_model.geom_size[geom_id] + # Per-axis scale: sphere is uniform (radius), ellipsoid uses its three radii. + is_ellipsoid = int(mj_model.geom_type[geom_id]) == mjtGeom.mjGEOM_ELLIPSOID + scale = np.asarray(size, dtype=np.float64) if is_ellipsoid else np.full(3, size[0]) + image = _cubemap_to_equirect(mj_model, texid, scale) + if image is None: + return None + + n_lat, n_lon = 24, 48 + lat = np.linspace(0.0, np.pi, n_lat + 1) + lon = np.linspace(0.0, 2.0 * np.pi, n_lon + 1) + lat_g, lon_g = np.meshgrid(lat, lon, indexing="ij") + sin_lat = np.sin(lat_g) + verts = np.stack( + [sin_lat * np.cos(lon_g), sin_lat * np.sin(lon_g), np.cos(lat_g)], axis=-1 + ).reshape(-1, 3) + # north pole (lat 0) -> top row of the image -> v=1 in glTF's v-up frame. + uv = np.stack([lon_g / (2.0 * np.pi), 1.0 - lat_g / np.pi], axis=-1).reshape(-1, 2) + + ncols = n_lon + 1 + i, j = np.meshgrid(np.arange(n_lat), np.arange(n_lon), indexing="ij") + a = (i * ncols + j).ravel() + b, c, d = a + 1, a + ncols + 1, a + ncols + # Wind triangles so their normals point outward (radius increases away + # from the sphere center), otherwise the sphere renders inside-out. + faces = np.concatenate( + [np.stack([a, c, b], axis=1), np.stack([a, d, c], axis=1)], axis=0 + ) + + mesh = trimesh.Trimesh(vertices=verts * scale, faces=faces, process=False) + material = trimesh.visual.material.PBRMaterial( + baseColorTexture=image, + metallicFactor=0.0, + roughnessFactor=1.0, + ) + mesh.visual = trimesh.visual.TextureVisuals(uv=uv, material=material) + return mesh + + +def _create_textured_plane_mesh( + mj_model: mujoco.MjModel, geom_id: int +) -> trimesh.Trimesh | None: + """Build a textured quad for a plane geom with a 2D material texture. + + Returns None when the plane has no material, no 2D texture, or the + texture is a cube map / unsupported format. UVs follow MuJoCo's + texrepeat/texuniform semantics so the image tiles across the plane. + The quad is double-sided so the plane stays visible from below. + """ + matid = int(mj_model.geom_matid[geom_id]) + if matid < 0 or matid >= mj_model.nmat: + return None + texid = _get_texture_id(mj_model, matid) + if texid < 0 or _is_cubemap_texture(mj_model, texid): + return None + image = _extract_texture_image(mj_model, texid) + if image is None: + return None + + size = mj_model.geom_size[geom_id] + half_x = float(size[0]) if size[0] > 0 else 10.0 + half_y = float(size[1]) if size[1] > 0 else 10.0 + + # texrepeat counts repetitions over the whole plane when texuniform is + # false, or repetitions per unit length when true. + rep = np.asarray(mj_model.mat_texrepeat[matid], dtype=np.float64) + if mj_model.mat_texuniform[matid]: + rep = rep * np.array([2.0 * half_x, 2.0 * half_y]) + + corners = np.array( + [ + [-half_x, -half_y, 0.0], + [half_x, -half_y, 0.0], + [half_x, half_y, 0.0], + [-half_x, half_y, 0.0], + ] + ) + corner_uvs = np.array([[0.0, 0.0], [rep[0], 0.0], [rep[0], rep[1]], [0.0, rep[1]]]) + # Duplicate the corners so the reversed faces get an outward -Z normal. + verts = np.vstack([corners, corners]) + uv = np.vstack([corner_uvs, corner_uvs]) + faces = np.array([[0, 1, 2], [0, 2, 3], [4, 6, 5], [4, 7, 6]], dtype=np.int64) + + mesh = trimesh.Trimesh(vertices=verts, faces=faces, process=False) + rgba = mj_model.mat_rgba[matid] + geom_rgba = mj_model.geom_rgba[geom_id] + # Blend when the material or geom is translucent, or the texture has alpha. + # Mirror the mesh texture path, which checks both rgba sources. + use_blending = rgba[-1] < 0.99 or geom_rgba[-1] < 0.99 or _has_alpha(image) + material = trimesh.visual.material.PBRMaterial( + baseColorFactor=rgba, + baseColorTexture=image, + metallicFactor=0.0, + roughnessFactor=1.0, + alphaMode="BLEND" if use_blending else "OPAQUE", + ) + mesh.visual = trimesh.visual.TextureVisuals(uv=uv, material=material) + return mesh + + +# Primitive geom types that take a cube map; planes take 2D textures. Used by +# both _textured_primitive_mesh and get_geom_texture_id so they stay in sync. +_CUBEMAP_PRIMITIVE_TYPES = ( + mjtGeom.mjGEOM_BOX, + mjtGeom.mjGEOM_SPHERE, + mjtGeom.mjGEOM_ELLIPSOID, +) + + +def _textured_primitive_mesh( + mj_model: mujoco.MjModel, geom_id: int +) -> trimesh.Trimesh | None: + """Build a textured mesh for a primitive geom, or None if not textured. + + Box/sphere/ellipsoid use the geom's cube map; planes use a 2D texture. + """ + geom_type = int(mj_model.geom_type[geom_id]) + if geom_type == mjtGeom.mjGEOM_BOX: + return _create_cubemap_box_mesh(mj_model, geom_id) + if geom_type in (mjtGeom.mjGEOM_SPHERE, mjtGeom.mjGEOM_ELLIPSOID): + return _create_textured_sphere_mesh(mj_model, geom_id) + if geom_type == mjtGeom.mjGEOM_PLANE: + return _create_textured_plane_mesh(mj_model, geom_id) + return None + + def _create_heightfield_mesh(mj_model: mujoco.MjModel, geom_id: int) -> trimesh.Trimesh: """Create a heightfield mesh, using the material texture when available.""" hfield_id = mj_model.geom_dataid[geom_id] @@ -500,10 +735,16 @@ def get_geom_texture_id(mj_model: mujoco.MjModel, geom_idx: int) -> int: if geom_type == mjtGeom.mjGEOM_HFIELD: return texid - # Box primitives use the cube map texture via per-face UV mapping. - if geom_type == mjtGeom.mjGEOM_BOX: + # Box/sphere/ellipsoid take a cube map (see _textured_primitive_mesh). + if geom_type in _CUBEMAP_PRIMITIVE_TYPES: return texid if _is_cubemap_texture(mj_model, texid) else -1 + # Planes take a 2D (non-cube-map) texture we can actually extract. + if geom_type == mjtGeom.mjGEOM_PLANE: + if _is_cubemap_texture(mj_model, texid): + return -1 + return texid if _is_2d_texture_supported(mj_model, texid) else -1 + if geom_type != mjtGeom.mjGEOM_MESH: return -1 @@ -519,13 +760,19 @@ def group_geoms_by_visual_compat( ) -> list[list[int]]: """Partition geom IDs into groups that can be safely merged. - Geoms sharing the same texture ID are grouped together. All - untextured geoms form a single group. + Geoms sharing the same texture ID are grouped together. Untextured + opaque geoms form one group; untextured translucent geoms are split by + color so each can be rendered with its own alpha-blended material. """ - groups: dict[int, list[int]] = {} + groups: dict[object, list[int]] = {} for gid in geom_ids: tex_id = get_geom_texture_id(mj_model, gid) - groups.setdefault(tex_id, []).append(gid) + if tex_id >= 0: + key: object = ("tex", tex_id) + else: + rgba = _resolve_flat_rgba(mj_model, gid) + key = ("opaque",) if rgba[3] >= 255 else ("alpha", tuple(int(c) for c in rgba)) + groups.setdefault(key, []).append(gid) return list(groups.values()) @@ -578,6 +825,41 @@ def _merge_meshes( return result +def _apply_translucent_blend(mesh: trimesh.Trimesh) -> None: + """Give a uniform-colored translucent mesh an alpha-blended material. + + trimesh exports per-vertex alpha as an opaque material, so a flat + translucent color (e.g. goal nets at rgba alpha 0.3) renders solid. + When every vertex shares one color with alpha < 1, swap to a PBR + material with alphaMode=BLEND so the viewer renders it see-through. + + This only handles a single translucent color because trimesh can't carry + per-vertex alpha plus a BLEND material. That is sufficient because + group_geoms_by_visual_compat splits untextured translucent geoms by exact + rgba, so every translucent merge reaching here is single-color + (test_translucent_geoms_split_by_color guards that invariant). + """ + vis = mesh.visual + if not isinstance(vis, trimesh.visual.ColorVisuals): + return + vc = vis.vertex_colors + if vc is None or len(vc) == 0: + return + unique = np.unique(vc, axis=0) + if len(unique) != 1 or unique[0, 3] >= 255: + return + mesh.visual = trimesh.visual.TextureVisuals( + uv=np.zeros((len(mesh.vertices), 2)), + material=trimesh.visual.material.PBRMaterial( + baseColorFactor=(unique[0] / 255.0).tolist(), + metallicFactor=0.0, + roughnessFactor=1.0, + alphaMode="BLEND", + doubleSided=True, + ), + ) + + def merge_geoms(mj_model: mujoco.MjModel, geom_ids: list[int]) -> trimesh.Trimesh: """Merge multiple geoms into a single trimesh in local body space.""" meshes = [] @@ -587,11 +869,13 @@ def merge_geoms(mj_model: mujoco.MjModel, geom_ids: list[int]) -> trimesh.Trimes else: meshes.append(create_primitive_mesh(mj_model, geom_id)) - return _merge_meshes( + result = _merge_meshes( meshes, [mj_model.geom_pos[gid] for gid in geom_ids], [mj_model.geom_quat[gid] for gid in geom_ids], ) + _apply_translucent_blend(result) + return result def _hull_trimesh_for_mesh_id( diff --git a/src/mjviser/scene.py b/src/mjviser/scene.py index 9923e8c..21f230c 100644 --- a/src/mjviser/scene.py +++ b/src/mjviser/scene.py @@ -16,6 +16,7 @@ from .conversions import ( get_body_name, + get_geom_texture_id, group_geoms_by_visual_compat, is_fixed_body, merge_geoms, @@ -1074,7 +1075,11 @@ def _add_fixed_geometry(self) -> None: body_id = self.mj_model.geom_bodyid[i] if not is_fixed_body(self.mj_model, body_id): continue - if self.mj_model.geom_type[i] == mjtGeom.mjGEOM_PLANE: + # Untextured planes render as an infinite reference grid. Textured + # planes fall through to the mesh path so their material shows. + if self.mj_model.geom_type[i] == mjtGeom.mjGEOM_PLANE and ( + get_geom_texture_id(self.mj_model, i) < 0 + ): body_name = get_body_name(self.mj_model, body_id) geom_name = mj_id2name(self.mj_model, mjtObj.mjOBJ_GEOM, i) self.server.scene.add_grid( diff --git a/tests/test_conversions.py b/tests/test_conversions.py index 70fef32..0c06ec7 100644 --- a/tests/test_conversions.py +++ b/tests/test_conversions.py @@ -10,10 +10,13 @@ from mjviser.conversions import ( _create_heightfield_mesh, _create_shape_mesh, + _cube_face_st, + _cubemap_sample_colors, _cubemap_vertex_colors, _extract_mesh_data, _extract_texture_image, _get_texture_id, + _gl_cube_face, _has_alpha, _merge_meshes, _resolve_flat_rgba, @@ -200,7 +203,8 @@ def test_extract_texture_flip(textured_model): """ -# MuJoCo cube face order with the outward axis each face sits on. +# MuJoCo cube face storage order (right, left, up, down, front, back) with the +# geom-local outward axis each face sits on (GL cube convention: up=+Y). _CUBE_FACE_AXES = ( (0, (1, 0, 0)), (1, (-1, 0, 0)), @@ -286,6 +290,268 @@ def test_get_geom_texture_id_cube_map_box(): assert get_geom_texture_id(model, 2) == -1 +# 2D textures on plane primitives + + +_TEXTURED_PLANE_XML = """ + + + + + + + + + + + + + + + + + +""" + + +def test_textured_plane_mesh_structure(): + # A plane with a 2D material texture becomes a flat, double-sided + # textured quad spanning the geom's full extent. + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + mesh = create_primitive_mesh(model, 0) + visual = _texture_visual(mesh) + assert mesh.faces.shape == (4, 3) + assert _pbr_material(mesh).baseColorTexture is not None + z = mesh.vertices[:, 2] + assert np.allclose(z, 0.0) + lo, hi = mesh.bounds + np.testing.assert_allclose(hi[:2], [5.0, 4.0]) + np.testing.assert_allclose(lo[:2], [-5.0, -4.0]) + assert visual.uv.shape == (8, 2) + + +def test_textured_plane_uv_repeat_non_uniform(): + # texuniform=false: texrepeat is repetitions across the whole plane. + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + uv = _texture_visual(create_primitive_mesh(model, 0)).uv + np.testing.assert_allclose(uv.max(axis=0), [2.0, 3.0]) + + +def test_textured_plane_uv_repeat_uniform(): + # texuniform=true: texrepeat is repetitions per unit length, so the + # max UV scales with the plane's full extent (0.5 * 10, 0.5 * 8). + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + uv = _texture_visual(create_primitive_mesh(model, 1)).uv + np.testing.assert_allclose(uv.max(axis=0), [5.0, 4.0]) + + +def test_plane_with_cube_map_falls_back_to_flat(): + # Cube maps are not supported on planes; they keep the flat fallback. + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + assert isinstance(create_primitive_mesh(model, 2).visual, trimesh.visual.ColorVisuals) + + +def test_untextured_planes_stay_flat(): + # Planes with only a flat material or no material keep the flat slab. + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + assert isinstance(create_primitive_mesh(model, 3).visual, trimesh.visual.ColorVisuals) + assert isinstance(create_primitive_mesh(model, 4).visual, trimesh.visual.ColorVisuals) + + +def test_get_geom_texture_id_textured_plane(): + # Textured planes group by their 2D texture; cube map and flat do not. + model = mujoco.MjModel.from_xml_string(_TEXTURED_PLANE_XML) + grid_texid = _get_texture_id(model, model.geom_matid[0]) + assert get_geom_texture_id(model, 0) == grid_texid + assert get_geom_texture_id(model, 2) == -1 + assert get_geom_texture_id(model, 3) == -1 + assert get_geom_texture_id(model, 4) == -1 + + +# Cube map textures on sphere / ellipsoid primitives + + +_CUBEMAP_SPHERE_XML = """ + + + + + + + + + + + + + + +""" + + +def test_textured_sphere_is_uv_textured(): + # A sphere with a cube map becomes a UV-textured mesh spanning the + # geom's radius, with a baked equirectangular image. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + mesh = create_primitive_mesh(model, 0) + assert isinstance(mesh.visual, trimesh.visual.TextureVisuals) + assert _pbr_material(mesh).baseColorTexture is not None + np.testing.assert_allclose(mesh.bounds[1], [0.3, 0.3, 0.3], atol=0.01) + + +def test_textured_sphere_normals_point_outward(): + # The lat-long sphere must wind outward, else it renders inside-out. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + mesh = create_primitive_mesh(model, 0) + centers = mesh.triangles_center + dirs = centers / np.maximum(np.linalg.norm(centers, axis=1, keepdims=True), 1e-9) + assert float(np.sum(mesh.face_normals * dirs, axis=1).mean()) > 0.5 + + +def test_textured_ellipsoid_extents(): + # An ellipsoid with a cube map spans its per-axis size. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + mesh = create_primitive_mesh(model, 1) + assert isinstance(mesh.visual, trimesh.visual.TextureVisuals) + np.testing.assert_allclose(mesh.bounds[1], [0.2, 0.3, 0.4], atol=0.02) + + +def test_equirect_bakes_face_colors(): + # Paint each cube face a distinct color; the baked equirectangular + # image should show each color along its axis direction. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + colors = np.array( + [ + [200, 0, 0], + [0, 200, 0], + [0, 0, 200], + [200, 200, 0], + [200, 0, 200], + [0, 200, 200], + ], + dtype=np.uint8, + ) + w = int(model.tex_width[0]) + adr = int(model.tex_adr[0]) + for i, color in enumerate(colors): + start = adr + i * w * w * 3 + model.tex_data[start : start + w * w * 3] = np.tile(color, w * w) + + # Geom cube convention: right, left, up, down, front, back -> +X, -X, +Y, + # -Y, +Z, -Z. + axes = np.array( + [[1, 0, 0], [-1, 0, 0], [0, 1, 0], [0, -1, 0], [0, 0, 1], [0, 0, -1]], + dtype=float, + ) + sampled = _cubemap_sample_colors(model, 0, axes) + assert sampled is not None + np.testing.assert_allclose(sampled[:, :3], colors, atol=2) + + +def test_sphere_without_cube_map_is_flat(): + # Plain spheres and spheres with a 2D texture keep the flat fallback. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + assert isinstance(create_primitive_mesh(model, 2).visual, trimesh.visual.ColorVisuals) + assert isinstance(create_primitive_mesh(model, 3).visual, trimesh.visual.ColorVisuals) + + +def test_get_geom_texture_id_textured_sphere(): + # Cube-mapped spheres and ellipsoids group by texture; others do not. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + assert get_geom_texture_id(model, 0) == 0 + assert get_geom_texture_id(model, 1) == 0 + assert get_geom_texture_id(model, 2) == -1 + assert get_geom_texture_id(model, 3) == -1 + + +def test_box_and_sphere_cube_orientation_consistent(): + # Box and sphere share _cube_face_st, so for any face-interior point the + # box's (s,t) for its known face equals what the sphere's face selection + # computes. Guards against the two paths drifting (cube corners are excluded + # since the major axis is ambiguous there). + rng = np.random.default_rng(0) + for fi, (major, sgn) in enumerate( + [(0, 1), (0, -1), (1, 1), (1, -1), (2, 1), (2, -1)] + ): + free = [k for k in range(3) if k != major] + for _ in range(20): + p = np.zeros(3) + p[major] = sgn + p[free[0]] = rng.uniform(-0.8, 0.8) + p[free[1]] = rng.uniform(-0.8, 0.8) + assert int(_gl_cube_face(p[None])[0]) == fi + s_box, t_box = _cube_face_st(p[None], np.array([fi])) + s_sph, t_sph = _cube_face_st(p[None], _gl_cube_face(p[None])) + assert abs(s_box[0] - s_sph[0]) < 1e-9 and abs(t_box[0] - t_sph[0]) < 1e-9 + + +def test_ellipsoid_samples_scaled_direction(): + # MuJoCo samples a geom cube texture by the scaled local position. For a flat + # ellipsoid, a +Z-leaning direction becomes +X-dominant after scaling, so it + # picks the right (+X) face, not front (+Z), unlike a sphere. + model = mujoco.MjModel.from_xml_string(_CUBEMAP_SPHERE_XML) + w = int(model.tex_width[0]) + adr = int(model.tex_adr[0]) + for face, color in {0: (255, 0, 0), 4: (0, 0, 255)}.items(): + model.tex_data[adr + face * w * w * 3 : adr + (face + 1) * w * w * 3] = np.tile( + color, w * w + ) + d = np.array([[0.5, 0.0, 0.8]]) + d = d / np.linalg.norm(d) + sphere_color = _cubemap_sample_colors(model, 0, d) + ell_color = _cubemap_sample_colors(model, 0, d * np.array([1.0, 1.0, 0.3])) + assert sphere_color is not None and ell_color is not None + np.testing.assert_array_equal(sphere_color[0, :3], (0, 0, 255)) # front + np.testing.assert_array_equal(ell_color[0, :3], (255, 0, 0)) # right + + +def test_textured_plane_geom_alpha_blends(): + # A textured plane made translucent via geom rgba (not material) blends. + xml = """ + + + + + + + + + +""" + model = mujoco.MjModel.from_xml_string(xml) + mesh = create_primitive_mesh(model, 0) + assert _pbr_material(mesh).alphaMode == "BLEND" + + +def test_merge_textured_spheres_keeps_texture(): + # Two spheres sharing a cube texture merge into one valid textured mesh. + xml = """ + + + + + + + + + + + + +""" + model = mujoco.MjModel.from_xml_string(xml) + mesh = merge_geoms(model, [0, 1]) + assert isinstance(mesh.visual, trimesh.visual.TextureVisuals) + assert mesh.visual.uv is not None + assert len(mesh.vertices) > 1000 + + # Mesh merging @@ -338,6 +604,47 @@ def test_group_geoms_by_visual_compat_splits_textured_hfield(): assert groups == [[0], [1]] +_TRANSLUCENT_XML = """ + + + + + + + + +""" + + +def test_translucent_geoms_split_by_color(): + # Opaque geoms group together; each translucent color is its own group. + model = mujoco.MjModel.from_xml_string(_TRANSLUCENT_XML) + groups = group_geoms_by_visual_compat(model, [0, 1, 2, 3]) + assert [0] in groups # opaque on its own (only solid is opaque here) + assert [1, 2] in groups # the two white nets share a group + assert [3] in groups # the red tint is separate + assert len(groups) == 3 + + +def test_translucent_merge_gets_blend_material(): + # A uniform translucent merge becomes a PBR BLEND material, not vertex + # colors (which trimesh would export as opaque). + model = mujoco.MjModel.from_xml_string(_TRANSLUCENT_XML) + mesh = merge_geoms(model, [1, 2]) + assert isinstance(mesh.visual, trimesh.visual.TextureVisuals) + mat = _pbr_material(mesh) + assert mat.alphaMode == "BLEND" + # trimesh stores baseColorFactor as uint8 internally; alpha < 255 => blended. + assert int(np.asarray(mat.baseColorFactor)[3]) < 255 + + +def test_opaque_merge_stays_vertex_colored(): + # Opaque geoms keep ColorVisuals (no spurious blend material). + model = mujoco.MjModel.from_xml_string(_TRANSLUCENT_XML) + mesh = merge_geoms(model, [0]) + assert isinstance(mesh.visual, trimesh.visual.ColorVisuals) + + def test_merge_geoms_hull(cubemap_model): mesh = merge_geoms_hull(cubemap_model, [0]) assert isinstance(mesh, trimesh.Trimesh)