google-deepmind / google-deepmind/mujoco

Support Exporting Flexes and Skins with the USDExporter

Open
#2,459 2 comments 0 reactions 0 assignees View on GitHub
enhancement
Dominant language
C++
Stars
15.2k
Forks
1.8k
Avg merge
10d 16h
Merged PRs (30d)
25

Description

### 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.

```python
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_

Contributor guide

Open the contributing guide

Assessment

This issue has not been assessed yet.

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.