Support Exporting Flexes and Skins with the USDExporter

[psclklnk]

The feature, motivation and pitch

Currently, the USDExporter does not support flexes and skins to be exported. Adding this functionality would allow for advanced rendering of deformable objects with MuJoCo and a rendering backend of the user's choice.

I had some luck with injecting a new USDObject subclass into the USDExporter when encountering skins in the MjvScene so I thought it may be beneficial to share it here.

class USDSkin(object_module.USDObject):

    def __init__(
        self,
        stage: Usd.Stage,
        model: mujoco.MjModel,
        scene: mujoco.MjvScene,
        geom: mujoco.MjvGeom,
        obj_name: str,
        dataid: int,
        rgba: np.ndarray = np.array([1, 1, 1, 1]),
        geom_textures: Sequence[Optional[Tuple[str, mujoco.mjtTexture]]] = (),
    ):
        super().__init__(stage, model, geom, obj_name, rgba, geom_textures)
        self.scene = scene

        self.dataid = dataid

        mesh_path = f"{self.xform_path}/Mesh_{obj_name}"
        self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
        self.usd_prim = stage.GetPrimAtPath(mesh_path)

        # setting mesh structure properties
        skin_vert, skin_face, skin_facenum = self._get_mesh_geometry()
        self.usd_mesh.GetPointsAttr().Set(skin_vert)
        self.usd_mesh.GetFaceVertexCountsAttr().Set([3 for _ in range(skin_facenum)])
        self.usd_mesh.GetFaceVertexIndicesAttr().Set(skin_face)

        if (
            geom.matid != -1
            and self.geom_textures[mujoco.mjtTextureRole.mjTEXROLE_RGB.value]
        ):
            # setting mesh uv properties
            mesh_texcoord, mesh_facetexcoord = self._get_uv_geometry()
            self.texcoords = UsdGeom.PrimvarsAPI(self.usd_mesh).CreatePrimvar(
                "UVMap",
                Sdf.ValueTypeNames.TexCoord2fArray,
                UsdGeom.Tokens.faceVarying,
            )
            self.texcoords.Set(mesh_texcoord)
            self.texcoords.SetIndices(Vt.IntArray(mesh_facetexcoord.tolist()))
            self.attach_image_material(self.usd_mesh)
        else:
            self.attach_solid_material(self.usd_mesh)

    def _get_facetexcoord_ranges(self, nmesh, arr):
        facetexcoords_ranges = [0]
        running_sum = 0
        for i in range(nmesh):
            running_sum += arr[i] * 3
            facetexcoords_ranges.append(running_sum)
        return facetexcoords_ranges

    def _get_uv_geometry(self):
        raise RuntimeError("UV Export not implemented")

    def _get_mesh_geometry(self):
        skin_vert_adr_from = self.scene.skinvertadr[self.dataid]
        skin_vert_adr_to = (
            self.scene.skinvertadr[self.dataid + 1]
            if self.dataid < self.model.nskin - 1
            else len(self.model.skin_vert)
        )
        skin_vert = self.scene.skinvert.reshape(-1, 3)[
            skin_vert_adr_from:skin_vert_adr_to
        ]

        skin_face_adr_from = self.model.skin_faceadr[self.dataid]
        skin_face_adr_to = (
            self.model.skin_faceadr[self.dataid + 1]
            if self.dataid < self.model.nskin - 1
            else len(self.model.skin_face)
        )
        skin_face = self.model.skin_face[skin_face_adr_from:skin_face_adr_to]
        assert np.min(skin_face) >= 0 and np.max(skin_face) < skin_vert.shape[0]
        skin_facenum = self.model.skin_facenum[self.dataid]
        assert skin_face.shape[0] == skin_facenum

        return skin_vert, skin_face, skin_facenum

    def update(
        self,
        pos: np.ndarray,
        mat: np.ndarray,
        visible: bool,
        frame: int,
        scale: Optional[np.ndarray] = None,
    ):
        if visible and frame - self.last_visible_frame > 1:
            # non consecutive visible frames
            self.update_visibility(False, max(0, self.last_visible_frame))
            self.update_visibility(True, frame)

        if visible:
            self.last_visible_frame = frame

        if scale is not None:
            self.update_scale(scale, frame)

        skin_vert = self._get_mesh_geometry()[0]
        self.usd_mesh.GetPointsAttr().Set(skin_vert, frame)

The class exports the skin vertex positions per frame. A major missing feature of this approach is that materials are currently not supported because I am not sure how those need to be accessed for skins. Furthermore, this class only targets skin objects and may not be the most efficient way of exporting the skin information per frame.

Alternatives

No response

Additional context

No response