chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions Packages Projects Releases Wiki Activity

Merge pull request #1714 from abhihjoshi:tendons

Browse Source 
PiperOrigin-RevId: 651048034
Change-Id: If5b7cb55ff0007c3c6f932e62f068d1168778a33
This commit is contained in:
Copybara-Service
2024-07-10 09:26:54 -07:00
parent 7f9f8b5429 e734527e7b
commit 60afda949c
9 changed files with 1092 additions and 969 deletions
Download Patch File Download Diff File Expand all files Collapse all files
.gitignore
+1
View File
@@ -18,6 +18,7 @@
*.egg-info/ *.egg-info/
build/ build/
build_cmake/ build_cmake/
python/dist/
# Exclude macOS folder attributes # Exclude macOS folder attributes
.DS_Store .DS_Store
python/mujoco/usd/camera.py
+53
View File
@@ -0,0 +1,53 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Camera handling for USD exporter."""
import mujoco.usd.utils as utils_component
import numpy as np
from pxr import Gf
from pxr import Usd
from pxr import UsdGeom
class USDCamera:
"""Class that handles the cameras in the USD scene."""
def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage
xform_path = f"/World/Camera_Xform_{obj_name}"
camera_path = f"{xform_path}/Camera_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_camera = UsdGeom.Camera.Define(stage, camera_path)
self.usd_prim = stage.GetPrimAtPath(camera_path)
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
self.usd_camera.CreateFocalLengthAttr().Set(12)
self.usd_camera.CreateFocusDistanceAttr().Set(400)
self.usd_camera.GetHorizontalApertureAttr().Set(12)
self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
"""Updates the position and orientation of the camera in the scene."""
transformation_mat = utils_component.create_transform_matrix(
rotation_matrix=cam_mat, translation_vector=cam_pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
python/mujoco/usd/component.py
-848
View File
@@ -1,848 +0,0 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
from typing import List, Optional, Tuple
import mujoco
import mujoco.usd.utils
import numpy as np
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
# pytype: disable=module-attr
from open3d import open3d as o3d
from pxr import Gf
from pxr import Sdf
from pxr import Usd
from pxr import UsdGeom
from pxr import UsdLux
from pxr import UsdShade
from pxr import Vt
class USDMesh:
def __init__(
self,
stage: Usd.Stage,
model: mujoco.MjModel,
geom: mujoco.MjvGeom,
obj_name: str,
dataid: int,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
self.stage = stage
self.model = model
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.dataid = dataid
self.texture_file = texture_file
xform_path = f"/World/Mesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/Mesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)
# setting mesh structure properties
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
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):
mesh_texcoord_adr_from = self.model.mesh_texcoordadr[self.dataid]
mesh_texcoord_adr_to = (
self.model.mesh_texcoordadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_texcoord)
)
mesh_texcoord = self.model.mesh_texcoord[
mesh_texcoord_adr_from:mesh_texcoord_adr_to
]
mesh_facetexcoord_ranges = self.get_facetexcoord_ranges(
self.model.nmesh, self.model.mesh_facenum
)
mesh_facetexcoord = self.model.mesh_facetexcoord.flatten()
mesh_facetexcoord = mesh_facetexcoord[
mesh_facetexcoord_ranges[self.dataid] : mesh_facetexcoord_ranges[
self.dataid + 1
]
]
mesh_facetexcoord[mesh_facetexcoord == len(mesh_texcoord)] = 0
return mesh_texcoord, mesh_facetexcoord
def _get_mesh_geometry(self):
mesh_vert_adr_from = self.model.mesh_vertadr[self.dataid]
mesh_vert_adr_to = (
self.model.mesh_vertadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_vert)
)
mesh_vert = self.model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
mesh_face_adr_from = self.model.mesh_faceadr[self.dataid]
mesh_face_adr_to = (
self.model.mesh_faceadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_face)
)
mesh_face = self.model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
mesh_facenum = self.model.mesh_facenum[self.dataid]
return mesh_vert, mesh_face, mesh_facenum
def _attach_material(self):
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
image_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Image_Texture")
)
uvmap_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("uvmap")
)
# setting the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput(
"diffuseColor", Sdf.ValueTypeNames.Color3f
).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file
)
image_shader.CreateInput(
"sourceColorSpace", Sdf.ValueTypeNames.Token
).Set("sRGB")
image_shader.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
uvmap_shader.ConnectableAPI(), "result"
)
image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
# setting uvmap shader attributes
uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
else:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
tuple(self.rgba[0:3])
)
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
transformation_mat = mujoco.usd.utils.create_transform_matrix(
rotation_matrix=mat, translation_vector=pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
self.update_visibility(visible, frame)
def update_visibility(self, visible: bool, frame: int):
if visible:
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
else:
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
class USDPrimitiveMesh:
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
self.stage = stage
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.texture_file = texture_file
self.usd_prim = Usd.Prim()
self.usd_mesh = UsdGeom.Mesh()
self.prim_mesh = None
self.transform_op = Gf.Matrix4d(1.)
def _set_refinement_properties(self):
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
def _get_uv_geometry(self):
assert self.prim_mesh
mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
# TODO(etom): bring back support for rescaling the texture coordinates.
return mesh_texcoord, mesh_facetexcoord.flatten()
def _get_mesh_geometry(self):
assert self.prim_mesh
# get mesh geometry from the open3d mesh model
mesh_vert = np.asarray(self.prim_mesh.vertices)
mesh_face = np.asarray(self.prim_mesh.triangles)
return mesh_vert, mesh_face, len(mesh_face)
def _attach_material(self):
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
image_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Image_Texture")
)
uvmap_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("uvmap")
)
# setting the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput(
"diffuseColor", Sdf.ValueTypeNames.Color3f
).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file
)
image_shader.CreateInput(
"sourceColorSpace", Sdf.ValueTypeNames.Token
).Set("sRGB")
image_shader.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
uvmap_shader.ConnectableAPI(), "result"
)
image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
# setting uvmap shader attributes
uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
else:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
tuple(self.rgba[:3])
)
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl)
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
transformation_mat = mujoco.usd.utils.create_transform_matrix(
rotation_matrix=mat, translation_vector=pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
self.update_visibility(visible, frame)
def update_visibility(self, visible: bool, frame: int):
if visible:
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
else:
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
class USDPrimitive:
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
self.stage = stage
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.texture_file = texture_file
self.usd_prim = Usd.Prim()
self.usd_primitive_shape = Usd.PrimitiveShape()
self.transform_op = Usd.TransformOp()
def _set_refinement_properties(self):
self.usd_prim.GetAttribute("subdivisionScheme").Set("none")
def _attach_material(self):
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
mtl = UsdShade.Material.Define(self.stage, mtl_path)
if self.texture_file:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
image_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Image_Texture")
)
uvmap_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("uvmap")
)
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput(
"diffuseColor", Sdf.ValueTypeNames.Color3f
).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file
)
image_shader.CreateInput(
"sourceColorSpace", Sdf.ValueTypeNames.Token
).Set("sRGB")
image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
uvmap_shader.ConnectableAPI(), "result"
)
image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
# setting uvmap shader attributes
uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
else:
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
tuple(self.rgba[:3])
)
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
self.usd_primitive_shape.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(self.usd_primitive_shape).Bind(mtl)
def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int):
transformation_mat = mujoco.usd_util.create_transform_matrix(
rotation_matrix=mat, translation_vector=pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
self.update_visibility(visible, frame)
def update_visibility(self, visible: bool, frame: int):
if visible:
self.usd_prim.GetAttribute("visibility").Set("inherited", frame)
else:
self.usd_prim.GetAttribute("visibility").Set("invisible", frame)
class USDCapsule(USDPrimitive):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/Capsule_Xform_{obj_name}"
capsule_path = f"{xform_path}/Capsule_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_primitive_shape = UsdGeom.Capsule.Define(stage, capsule_path)
self.usd_prim = stage.GetPrimAtPath(capsule_path)
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
self.scale_op = self.usd_xform.AddScaleOp()
# setting attributes for the shape
self._set_size_attributes()
self._attach_material()
self._set_refinement_properties()
def _set_size_attributes(self):
self.usd_primitive_shape.GetRadiusAttr().Set(float(self.geom.size[0]))
self.usd_primitive_shape.GetHeightAttr().Set(
float(self.geom.size[2] * 2)
) # mujoco gives the half length
class USDEllipsoid(USDPrimitive):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/Ellipsoid_Xform_{obj_name}"
ellipsoid_path = f"{xform_path}/Ellipsoid_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_primitive_shape = UsdGeom.Sphere.Define(stage, ellipsoid_path)
self.usd_prim = stage.GetPrimAtPath(ellipsoid_path)
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
self.scale_op = self.usd_xform.AddScaleOp()
# setting attributes for the shape
self._set_size_attributes()
self._attach_material()
self._set_refinement_properties()
def _set_size_attributes(self):
self.scale_op.Set(Gf.Vec3d(self.geom.size.tolist()))
class USDCubeMesh(USDPrimitiveMesh):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/CubeMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/CubeMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)
self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
width=self.geom.size[0] * 2,
height=self.geom.size[1] * 2,
depth=self.geom.size[2] * 2,
)
self.prim_mesh.translate(-self.prim_mesh.get_center())
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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([i for i in range(mesh_facenum * 3)]))
self._set_refinement_properties()
# setting attributes for the shape
self._attach_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
class USDSphereMesh(USDPrimitiveMesh):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/SphereMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/SphereMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)
self.prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
radius=float(self.geom.size[0]), create_uv_map=True
)
self.prim_mesh.translate(-self.prim_mesh.get_center())
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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([i for i in range(mesh_facenum * 3)]))
self._set_refinement_properties()
# setting attributes for the shape
self._attach_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
class USDCylinderMesh(USDPrimitiveMesh):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/CylinderMesh_Xform_{obj_name}"
mesh_path = f"{xform_path}/CylinderMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_prim = stage.GetPrimAtPath(mesh_path)
self.prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
radius=self.geom.size[0],
height=self.geom.size[2] * 2,
create_uv_map=True,
)
self.prim_mesh.translate(-self.prim_mesh.get_center())
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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([i for i in range(mesh_facenum * 3)]))
self._set_refinement_properties()
# setting attributes for the shape
self._attach_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
class USDPlaneMesh(USDPrimitiveMesh):
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
xform_path = f"/World/Plane_Xform_{obj_name}"
plane_path = f"{xform_path}/PlaneMesh_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_mesh = UsdGeom.Mesh.Define(stage, plane_path)
self.usd_prim = stage.GetPrimAtPath(plane_path)
self.prim_mesh = o3d.geometry.TriangleMesh.create_box(
width=self.geom.size[0] * 2 if self.geom.size[0] > 0 else 100,
height=self.geom.size[1] * 2 if self.geom.size[1] > 0 else 100,
depth=0.001,
create_uv_map=True,
map_texture_to_each_face=True,
)
self.prim_mesh.translate(-self.prim_mesh.get_center())
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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([i for i in range(mesh_facenum * 3)]))
self._set_refinement_properties()
# setting attributes for the shape
self._attach_material()
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
class USDSphereLight:
def __init__(
self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3
):
self.stage = stage
xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.SphereLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path)
# we assume in mujoco that all lights are point lights
self.usd_light.GetRadiusAttr().Set(radius)
self.usd_light.GetTreatAsPointAttr().Set(False)
self.usd_light.GetNormalizeAttr().Set(True)
# defining ops required by update function
self.translate_op = self.usd_xform.AddTranslateOp()
def update(
self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int
):
self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
if not np.any(pos):
intensity = 0
self.usd_light.GetIntensityAttr().Set(intensity)
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
class USDDomeLight:
def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage
xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.DomeLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path)
# we assume in mujoco that all lights are point lights
self.usd_light.GetNormalizeAttr().Set(True)
def update(self, intensity: int, color: np.ndarray, frame: int):
self.usd_light.GetIntensityAttr().Set(intensity)
self.usd_light.GetExposureAttr().Set(0.0)
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
class USDCamera:
def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage
xform_path = f"/World/Camera_Xform_{obj_name}"
camera_path = f"{xform_path}/Camera_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_camera = UsdGeom.Camera.Define(stage, camera_path)
self.usd_prim = stage.GetPrimAtPath(camera_path)
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
# self.usd_camera.CreateFocalLengthAttr().Set(18.14756) # default in omniverse
self.usd_camera.CreateFocalLengthAttr().Set(12)
self.usd_camera.CreateFocusDistanceAttr().Set(400)
self.usd_camera.GetHorizontalApertureAttr().Set(12)
self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int):
transformation_mat = mujoco.usd_util.create_transform_matrix(
rotation_matrix=cam_mat, translation_vector=cam_pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
python/mujoco/usd/demo.py
+89
View File
@@ -0,0 +1,89 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Demo script for USD exporter."""
import argparse
import pathlib
import mujoco
from mujoco.usd import exporter
def generate_usd_trajectory(local_args):
"""Generates a USD file given the user arguments."""
# load a model to mujoco
model_path = local_args.model_path
m = mujoco.MjModel.from_xml_path(model_path)
d = mujoco.MjData(m)
# create an instance of the USDExporter
exp = exporter.USDExporter(
model=m,
output_directory_name=pathlib.Path(local_args.model_path).stem,
output_directory_root=local_args.output_directory_root,
camera_names=local_args.camera_names,
)
# step through the simulation for the given duration of time
while d.time < local_args.duration:
mujoco.mj_step(m, d)
if exp.frame_count < d.time * local_args.framerate:
exp.update_scene(data=d)
exp.save_scene(filetype=local_args.export_extension)
if __name__ == '__main__':
parser = argparse.ArgumentParser()
parser.add_argument(
'--model_path', type=str, required=True, help='path to mjcf xml model'
)
parser.add_argument(
'--duration',
type=int,
default=5,
help='duration in seconds for the generated video',
)
parser.add_argument(
'--framerate',
type=int,
default=60,
help='frame rate of the generated video',
)
parser.add_argument(
'--output_directory_root',
type=str,
default='../usd_trajectories/',
help='location where to create usd files',
)
parser.add_argument(
'--camera_names', type=str, nargs='+', help='cameras to include in usd'
)
parser.add_argument(
'--export_extension',
type=str,
default='usd',
help='extension of exported file (usd, usda, usdc)',
)
args = parser.parse_args()
generate_usd_trajectory(args)
python/mujoco/usd/exporter.py
+167 -109
View File
@@ -12,27 +12,30 @@
# See the License for the specific language governing permissions and # See the License for the specific language governing permissions and
# limitations under the License. # limitations under the License.
# ============================================================================== # ==============================================================================
"""USD exporter."""
import os import os
from typing import List, Optional
import mujoco import mujoco
import mujoco.usd.component as component_module import mujoco.usd.camera as camera_module
import mujoco.usd.lights as light_module
import mujoco.usd.objects as object_module
import mujoco.usd.shapes as shapes_module
import numpy as np import numpy as np
from PIL import Image as im
from PIL import ImageOps
import scipy import scipy
import termcolor import termcolor
import tqdm import tqdm
from typing import List, Optional, Tuple, Union
from PIL import Image as im
from PIL import ImageOps
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
# pytype: disable=module-attr
from pxr import Sdf from pxr import Sdf
from pxr import Usd from pxr import Usd
from pxr import UsdGeom from pxr import UsdGeom
class USDExporter: class USDExporter:
"""MuJoCo to USD exporter for porting scenes to external renderers."""
def __init__( def __init__(
self, self,
@@ -47,13 +50,13 @@ class USDExporter:
specialized_materials_file: Optional[str] = None, specialized_materials_file: Optional[str] = None,
verbose: bool = True, verbose: bool = True,
): ):
"""Initializes a new USD Exporter """Initializes a new USD Exporter.
Args: Args:
model: an MjModel instance. model: an MjModel instance.
height: image height in pixels. height: image height in pixels.
width: image width in pixels. width: image width in pixels.
max_geom: Optional integer specifying the maximum number of geoms that max_geom: optional integer specifying the maximum number of geoms that
can be rendered in the same scene. If None this will be chosen can be rendered in the same scene. If None this will be chosen
automatically based on the estimated maximum number of renderable automatically based on the estimated maximum number of renderable
geoms in the model. geoms in the model.
@@ -61,6 +64,10 @@ class USDExporter:
and assets generated by the USD renderer. and assets generated by the USD renderer.
output_directory_root: path to root directory storing generated frames output_directory_root: path to root directory storing generated frames
and assets by the USD renderer. and assets by the USD renderer.
light_intensity: intensity of the light in the scene.
camera_names: list of camera names to be used in the scene.
specialized_materials_file: path to a file containing a list of
materials to be used in the scene.
verbose: decides whether to print updates. verbose: decides whether to print updates.
""" """
@@ -99,7 +106,12 @@ class USDExporter:
self.frame_count = 0 # maintains how many times we have saved the scene self.frame_count = 0 # maintains how many times we have saved the scene
self.updates = 0 self.updates = 0
self.geom_name2usd = {} self.geom_names = set()
self.geom_refs = {}
# initializing list of lights and cameras
self.usd_lights = []
self.usd_cameras = []
# initializing rendering requirements # initializing rendering requirements
self.renderer = mujoco.Renderer(model, height, width, max_geom) self.renderer = mujoco.Renderer(model, height, width, max_geom)
@@ -114,18 +126,21 @@ class USDExporter:
@property @property
def usd(self): def usd(self):
"""Returns the USD file as a string."""
return self.stage.GetRootLayer().ExportToString() return self.stage.GetRootLayer().ExportToString()
@property @property
def scene(self): def scene(self):
"""Returns the scene."""
return self.renderer.scene return self.renderer.scene
def _initialize_usd_stage(self): def _initialize_usd_stage(self):
"""Initializes a USD stage to represent the mujoco scene."""
self.stage = Usd.Stage.CreateInMemory() self.stage = Usd.Stage.CreateInMemory()
UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z) UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
self.stage.SetStartTimeCode(0) self.stage.SetStartTimeCode(0)
# add as user input # add as user input
self.stage.SetTimeCodesPerSecond(24.0) self.stage.SetTimeCodesPerSecond(60.0)
default_prim = UsdGeom.Xform.Define( default_prim = UsdGeom.Xform.Define(
self.stage, Sdf.Path("/World") self.stage, Sdf.Path("/World")
@@ -133,6 +148,7 @@ class USDExporter:
self.stage.SetDefaultPrim(default_prim) self.stage.SetDefaultPrim(default_prim)
def _initialize_output_directories(self): def _initialize_output_directories(self):
"""Initializes output directories to store frames and assets."""
self.output_directory_path = os.path.join( self.output_directory_path = os.path.join(
self.output_directory_root, self.output_directory_name self.output_directory_root, self.output_directory_name
) )
@@ -161,7 +177,7 @@ class USDExporter:
data: mujoco.MjData, data: mujoco.MjData,
scene_option: Optional[mujoco.MjvOption] = None, scene_option: Optional[mujoco.MjvOption] = None,
): ):
"""Updates the scene with latest sim data """Updates the scene with latest sim data.
Args: Args:
data: structure storing current simulation state data: structure storing current simulation state
@@ -175,7 +191,6 @@ class USDExporter:
# update the mujoco renderer # update the mujoco renderer
self.renderer.update_scene(data, scene_option=scene_option) self.renderer.update_scene(data, scene_option=scene_option)
# TODO: update scene options
if self.updates == 0: if self.updates == 0:
self._initialize_usd_stage() self._initialize_usd_stage()
@@ -189,7 +204,7 @@ class USDExporter:
self.updates += 1 self.updates += 1
def _load_textures(self): def _load_textures(self):
# TODO: remove code once added internally to mujoco """Load textures."""
data_adr = 0 data_adr = 0
self.texture_files = [] self.texture_files = []
for texture_id in tqdm.tqdm(range(self.model.ntex)): for texture_id in tqdm.tqdm(range(self.model.ntex)):
@@ -212,7 +227,7 @@ class USDExporter:
) )
img_path = os.path.join( img_path = os.path.join(
relative_path, texture_file_name relative_path, texture_file_name
) # relative path, TODO: switch back to this )
self.texture_files.append(img_path) self.texture_files.append(img_path)
@@ -228,19 +243,20 @@ class USDExporter:
) )
def _load_geom(self, geom: mujoco.MjvGeom): def _load_geom(self, geom: mujoco.MjvGeom):
"""Loads a geom into the USD scene."""
geom_name = self._get_geom_name(geom)
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid) assert geom_name not in self.geom_names
assert geom_name not in self.geom_name2usd
texture_file = ( texture_file = (
self.texture_files[self.model.mat_texid[mujoco.mjNTEXMAT*geom.matid]] self.texture_files[self.model.mat_texid[geom.matid][0]]
if geom.matid != -1 if geom.matid != -1
else None else None
) )
# handles meshes in scene # handling meshes in our scene
if geom.type == mujoco.mjtGeom.mjGEOM_MESH: if geom.type == mujoco.mjtGeom.mjGEOM_MESH:
usd_geom = component_module.USDMesh( usd_geom = object_module.USDMesh(
stage=self.stage, stage=self.stage,
model=self.model, model=self.model,
geom=geom, geom=geom,
@@ -249,95 +265,77 @@ class USDExporter:
rgba=geom.rgba, rgba=geom.rgba,
texture_file=texture_file, texture_file=texture_file,
) )
elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
usd_geom = component_module.USDPlaneMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
usd_geom = component_module.USDSphereMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
usd_geom = component_module.USDCapsule(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
usd_geom = component_module.USDEllipsoid(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
usd_geom = component_module.USDCylinderMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
usd_geom = component_module.USDCubeMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
else: else:
usd_geom = None # handling tendons in our scene
if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
mesh_config = shapes_module.mesh_config_generator(
name=geom_name,
geom_type=geom.type,
size=np.array([1.0, 1.0, 1.0]),
decouple=True
)
usd_geom = object_module.USDTendon(
mesh_config=mesh_config,
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
# handling primitives in our scene
else:
mesh_config = shapes_module.mesh_config_generator(
name=geom_name,
geom_type=geom.type,
size=geom.size
)
usd_geom = object_module.USDPrimitiveMesh(
mesh_config=mesh_config,
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
self.geom_name2usd[geom_name] = usd_geom self.geom_names.add(geom_name)
self.geom_refs[geom_name] = usd_geom
def _update_geoms(self): def _update_geoms(self):
"""Iterate through all geoms in the scene and makes update."""
geom_names = set(self.geom_name2usd.keys())
# iterate through all geoms in the scene and makes update
for i in range(self.scene.ngeom): for i in range(self.scene.ngeom):
geom = self.scene.geoms[i] geom = self.scene.geoms[i]
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid) geom_name = self._get_geom_name(geom)
if geom_name not in self.geom_name2usd: if geom_name not in self.geom_names:
# load a new object into USD
self._load_geom(geom) self._load_geom(geom)
if self.geom_name2usd[geom_name]:
self.geom_name2usd[geom_name].update_visibility(False, 0)
if self.geom_name2usd[geom_name]: if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
self.geom_name2usd[geom_name].update( tendon_scale = geom.size
self.geom_refs[geom_name].update(
pos=geom.pos,
mat=geom.mat,
scale=tendon_scale,
visible=geom.rgba[3] > 0,
frame=self.updates,
)
else:
self.geom_refs[geom_name].update(
pos=geom.pos, pos=geom.pos,
mat=geom.mat, mat=geom.mat,
visible=geom.rgba[3] > 0, visible=geom.rgba[3] > 0,
frame=self.updates, frame=self.updates,
) )
if geom_name in geom_names:
geom_names.remove(geom_name)
for geom_name in geom_names:
if self.geom_name2usd[geom_name]:
self.geom_name2usd[geom_name].update_visibility(False, self.updates)
def _load_lights(self): def _load_lights(self):
# initializes an usd light object for every light in the scene # initializes an usd light object for every light in the scene
self.usd_lights = []
for i in range(self.scene.nlight): for i in range(self.scene.nlight):
light = self.scene.lights[i] light = self.scene.lights[i]
if not np.allclose(light.pos, [0, 0, 0]): if not np.allclose(light.pos, [0, 0, 0]):
self.usd_lights.append self.usd_lights.append(
(component_module.USDSphereLight(stage=self.stage, obj_name=str(i))) light_module.USDSphereLight(stage=self.stage, obj_name=str(i))
)
else: else:
self.usd_lights.append(None) self.usd_lights.append(None)
@@ -359,19 +357,23 @@ class USDExporter:
) )
def _load_cameras(self): def _load_cameras(self):
self.usd_cameras = []
if self.camera_names is not None: if self.camera_names is not None:
for name in self.camera_names: for name in self.camera_names:
self.usd_cameras.append( self.usd_cameras.append(
component_module.USDCamera(stage=self.stage, obj_name=name)) camera_module.USDCamera(stage=self.stage, obj_name=name))
def _update_cameras( def _update_cameras(
self, self,
data: mujoco.MjData, data: mujoco.MjData,
scene_option: Optional[mujoco.MjvOption] = None, scene_option: Optional[mujoco.MjvOption] = None,
): ):
for i in range(len(self.usd_cameras)): """Updates cameras.
Args:
data: An MjData instance.
scene_option: An optional MjvOption instance.
"""
for i in range(len(self.usd_cameras)):
camera = self.usd_cameras[i] camera = self.usd_cameras[i]
camera_name = self.camera_names[i] camera_name = self.camera_names[i]
@@ -386,12 +388,14 @@ class USDExporter:
up = avg_camera.up up = avg_camera.up
right = np.cross(forward, up) right = np.cross(forward, up)
R = np.eye(3) rotation = np.eye(3)
R[:, 0] = right rotation[:, 0] = right
R[:, 1] = up rotation[:, 1] = up
R[:, 2] = -forward rotation[:, 2] = -forward
camera.update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates) camera.update(
cam_pos=avg_camera.pos, cam_mat=rotation, frame=self.updates
)
def add_light( def add_light(
self, self,
@@ -399,37 +403,91 @@ class USDExporter:
intensity: int, intensity: int,
radius: Optional[float] = 1.0, radius: Optional[float] = 1.0,
color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]), color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]),
objid: Optional[int] = 1, obj_name: Optional[str] = "light_1",
light_type: Optional[str] = "sphere", light_type: Optional[str] = "sphere",
): ):
"""Adds a user defined, fixed light.
Args:
pos: position of the light in 3D space.
intensity: intensity of the light.
radius: radius of the light to be used by renderer.
color: color of the light.
obj_name: name associated with the light.
light_type: type of light (sphere or dome).
"""
if light_type == "sphere": if light_type == "sphere":
new_light = component_module.USDSphereLight(stage=self.stage, obj_name=str(objid), radius=radius) new_light = light_module.USDSphereLight(
stage=self.stage, obj_name=obj_name, radius=radius
)
new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0) new_light.update(
pos=np.array(pos), intensity=intensity, color=color, frame=0
)
elif light_type == "dome": elif light_type == "dome":
new_light = component_module.USDDomeLight( new_light = light_module.USDDomeLight(stage=self.stage, obj_name=obj_name)
stage=self.stage, obj_name=str(objid)) new_light.update(intensity=intensity, color=color)
new_light.update(intensity=intensity, color=color, frame=0)
def add_camera( def add_camera(
self, self,
pos: List[float], pos: List[float],
rotation_xyz: List[float], rotation_xyz: List[float],
objid: Optional[int] = 1, obj_name: Optional[str] = "camera_1",
): ):
new_camera = component_module.USDCamera( """Adds a user defined, fixed camera.
stage=self.stage, obj_name=str(objid))
Args:
pos: position of the camera in 3D space.
rotation_xyz: euler rotation of the camera.
obj_name: name associated with the camera.
"""
new_camera = camera_module.USDCamera(
stage=self.stage, obj_name=obj_name)
r = scipy.spatial.transform.Rotation.from_euler( r = scipy.spatial.transform.Rotation.from_euler(
"xyz", rotation_xyz, degrees=True) "xyz", rotation_xyz, degrees=True)
new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0) new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
def save_scene(self, filetype: str = "usd"): def save_scene(self, filetype: str = "usd"):
"""Saves the scene to a USD file."""
assert filetype in ["usd", "usda", "usdc"]
self.stage.SetEndTimeCode(self.frame_count) self.stage.SetEndTimeCode(self.frame_count)
# post-processing for visibility of geoms in scene
for _, geom_ref in self.geom_refs.items():
geom_ref.update_visibility(False, geom_ref.last_visible_frame+1)
self.stage.Export( self.stage.Export(
f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}" f"{self.output_directory_root}/{self.output_directory_name}/" +
f"frames/frame_{self.frame_count}.{filetype}"
) )
if self.verbose: if self.verbose:
print(termcolor.colored(f"Writing frame_{self.frame_count}", "green")) print(
termcolor.colored(
f"Completed writing frame_{self.frame_count}.{filetype}", "green"
)
)
def _get_geom_name(self, geom):
"""Adding id as part of name for USD file."""
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
if not geom_name:
geom_name = "None"
geom_name += f"_id{geom.objid}"
# adding additional naming information to differentiate
# between geoms and tendons
if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM:
geom_name += "_geom"
elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
geom_name += f"_tendon_segid{geom.segid}"
return geom_name
# for debugging purposes, prints all geoms in scene
# including those part of tendons
def _print_scene_geom_info(self):
for i in range(self.scene.ngeom):
geom = self.scene.geoms[i]
geom_name = self._get_geom_name(geom)
print(i, geom_name)
python/mujoco/usd/exporter_test.py
+10 -12
View File
@@ -17,21 +17,19 @@
import logging import logging
import os import os
import mujoco
from absl.testing import absltest from absl.testing import absltest
from etils import epath from etils import epath
import mujoco
# Open3D and USD are not fully supported on all MuJoCo architectures. # Open3D and USD are not fully supported on all MuJoCo architectures.
# pylint: disable=python.style(g-import-not-at-top)
execute_test = True execute_test = True
try: try:
from mujoco.usd import exporter as exporter_module from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top
from pxr import Usd
except ImportError: except ImportError:
logging.warning('Skipping test due to missing import') logging.warning('Skipping test due to missing import')
execute_test = False execute_test = False
# pylint: enable=python.style(g-import-not-at-top)
class ExporterTest(absltest.TestCase): class ExporterTest(absltest.TestCase):
@@ -52,19 +50,19 @@ class ExporterTest(absltest.TestCase):
data = mujoco.MjData(model) data = mujoco.MjData(model)
exporter = exporter_module.USDExporter( exporter = exporter_module.USDExporter(
model, model,
output_directory_name="mujoco_usdpkg", output_directory_name='mujoco_usdpkg',
output_directory_root=output_dir, output_directory_root=output_dir,
) )
exporter.update_scene(data) exporter.update_scene(data)
exporter.save_scene("export.usda") exporter.save_scene('export.usda')
with open(os.path.join( with open(os.path.join(
output_dir, "mujoco_usdpkg/frames", "frame_1_.export.usda"), "r") as f: output_dir, 'mujoco_usdpkg/frames', 'frame_1.export.usda'), 'r') as f:
golden_path = os.path.join( golden_path = os.path.join(
epath.resource_path("mujoco"), "testdata", "usd_golden.usda") epath.resource_path('mujoco'), 'testdata', 'usd_golden.usda')
with open(golden_path, "r") as golden_file: with open(golden_path, 'r') as golden_file:
self.assertEqual(f.read(), golden_file.read()) self.assertEqual(f.read(), golden_file.read())
if __name__ == "__main__": if __name__ == '__main__':
absltest.main() absltest.main()
python/mujoco/usd/lights.py
+81
View File
@@ -0,0 +1,81 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Light handling for USD exporter."""
from typing import Optional
import numpy as np
from pxr import Gf
from pxr import Usd
from pxr import UsdGeom
from pxr import UsdLux
class USDSphereLight:
"""Class that handles the sphere lights in the USD scene."""
def __init__(
self, stage: Usd.Stage, obj_name: str, radius: Optional[float] = 0.3
):
self.stage = stage
xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.SphereLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path)
# we assume in mujoco that all lights are point lights
self.usd_light.GetRadiusAttr().Set(radius)
self.usd_light.GetTreatAsPointAttr().Set(False)
self.usd_light.GetNormalizeAttr().Set(True)
# defining ops required by update function
self.translate_op = self.usd_xform.AddTranslateOp()
def update(
self, pos: np.ndarray, intensity: int, color: np.ndarray, frame: int
):
"""Updates the attributes of a sphere light."""
self.translate_op.Set(Gf.Vec3d(pos.tolist()), frame)
if not np.any(pos):
intensity = 0
self.usd_light.GetIntensityAttr().Set(intensity)
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
class USDDomeLight:
"""Class that handles the dome lights in the USD scene."""
def __init__(self, stage: Usd.Stage, obj_name: str):
self.stage = stage
xform_path = f"/World/Light_Xform_{obj_name}"
light_path = f"{xform_path}/Light_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, xform_path)
self.usd_light = UsdLux.DomeLight.Define(stage, light_path)
self.usd_prim = stage.GetPrimAtPath(light_path)
# we assume in mujoco that all lights are point lights
self.usd_light.GetNormalizeAttr().Set(True)
def update(self, intensity: int, color: np.ndarray):
"""Updates the attributes of a dome light."""
self.usd_light.GetIntensityAttr().Set(intensity)
self.usd_light.GetExposureAttr().Set(0.0)
self.usd_light.GetColorAttr().Set(Gf.Vec3d(color.tolist()))
python/mujoco/usd/objects.py
+500
View File
@@ -0,0 +1,500 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Objects module for USD exporter."""
import abc
import collections
from typing import Optional
import mujoco
import mujoco.usd.shapes as shapes_component
import mujoco.usd.utils as utils_component
import numpy as np
from pxr import Gf
from pxr import Sdf
from pxr import Usd
from pxr import UsdGeom
from pxr import UsdShade
from pxr import Vt
class USDObject(abc.ABC):
"""Abstract interface for all USD objects including meshes and primitives.
Subclasses must implement:
* `_get_uv_geometry(self)`: gets the nessecary UV information to
wrap a texture around an object in USD. Each subclass implements
their own method to getting UV information as different objects
are contructed in different ways.
* `_get_mesh_geometry(self)`: gets the mesh geometry of an object
in the scene.
"""
def __init__(
self,
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
self.stage = stage
self.geom = geom
self.obj_name = obj_name
self.rgba = rgba
self.texture_file = texture_file
self.xform_path = f"/World/Mesh_Xform_{obj_name}"
self.usd_xform = UsdGeom.Xform.Define(stage, self.xform_path)
# defining ops required by update function
self.transform_op = self.usd_xform.AddTransformOp()
self.scale_op = self.usd_xform.AddScaleOp()
self.last_visible_frame = -2
@abc.abstractmethod
def _get_uv_geometry(self):
"""Gets UV information for an object in the scene."""
raise NotImplementedError
@abc.abstractmethod
def _get_mesh_geometry(self):
"""Gets structure of an object in the scene."""
raise NotImplementedError
def attach_image_material(self, usd_mesh):
"""Attaches an image texture to a material for a USD object."""
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
mtl = UsdShade.Material.Define(self.stage, mtl_path)
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
image_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Image_Texture")
)
uvmap_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("uvmap")
)
# setting the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput(
"diffuseColor", Sdf.ValueTypeNames.Color3f
).ConnectToSource(image_shader.ConnectableAPI(), "rgb")
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
# setting the image texture attributes
image_shader.CreateIdAttr("UsdUVTexture")
image_shader.CreateInput("file", Sdf.ValueTypeNames.Asset).Set(
self.texture_file
)
image_shader.CreateInput("sourceColorSpace", Sdf.ValueTypeNames.Token).Set(
"sRGB"
)
image_shader.CreateInput("wrapS", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("wrapT", Sdf.ValueTypeNames.Token).Set("repeat")
image_shader.CreateInput("st", Sdf.ValueTypeNames.Float2).ConnectToSource(
uvmap_shader.ConnectableAPI(), "result"
)
image_shader.CreateOutput("rgb", Sdf.ValueTypeNames.Float3)
# setting uvmap shader attributes
uvmap_shader.CreateIdAttr("UsdPrimvarReader_float2")
uvmap_shader.CreateInput("varname", Sdf.ValueTypeNames.Token).Set("UVMap")
uvmap_shader.CreateOutput("results", Sdf.ValueTypeNames.Float2)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
def attach_solid_material(self, usd_mesh):
"""Attaches an solid texture to a material for a USD object."""
mtl_path = Sdf.Path(f"/World/_materials/Material_{self.obj_name}")
mtl = UsdShade.Material.Define(self.stage, mtl_path)
bsdf_shader = UsdShade.Shader.Define(
self.stage, mtl_path.AppendPath("Principled_BSDF")
)
# settings the bsdf shader attributes
bsdf_shader.CreateIdAttr("UsdPreviewSurface")
bsdf_shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(
tuple(self.rgba[0:3])
)
bsdf_shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(
float(self.rgba[-1])
)
bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(
self.geom.shininess
)
bsdf_shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(
1.0 - self.geom.shininess
)
mtl.CreateSurfaceOutput().ConnectToSource(
bsdf_shader.ConnectableAPI(), "surface"
)
usd_mesh.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
UsdShade.MaterialBindingAPI(usd_mesh).Bind(mtl)
def _set_refinement_properties(self, usd_prim, scheme="none"):
usd_prim.GetAttribute("subdivisionScheme").Set(scheme)
def update(
self,
pos: np.ndarray,
mat: np.ndarray,
visible: bool,
frame: int,
scale: Optional[np.ndarray] = None,
):
"""Updates the position and orientation of an object."""
transformation_mat = utils_component.create_transform_matrix(
rotation_matrix=mat, translation_vector=pos
).T
self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
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)
def update_visibility(self, visible: bool, frame: int):
"""Updates the visibility of an object in a scene for a given frame."""
visibility_setting = "inherited" if visible else "invisible"
self.usd_xform.GetVisibilityAttr().Set(visibility_setting, frame)
def update_scale(self, scale: np.ndarray, frame: int):
"""Updates the scale of an object in the scene for a given frame."""
self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame)
class USDMesh(USDObject):
"""Class that handles predefined meshes in the USD scene."""
def __init__(
self,
stage: Usd.Stage,
model: mujoco.MjModel,
geom: mujoco.MjvGeom,
obj_name: str,
dataid: int,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
self.model = model
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
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# 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()))
if self.texture_file:
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):
mesh_texcoord_adr_from = self.model.mesh_texcoordadr[self.dataid]
mesh_texcoord_adr_to = (
self.model.mesh_texcoordadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_texcoord)
)
mesh_texcoord = self.model.mesh_texcoord[
mesh_texcoord_adr_from:mesh_texcoord_adr_to
]
mesh_facetexcoord_ranges = self._get_facetexcoord_ranges(
self.model.nmesh, self.model.mesh_facenum
)
mesh_facetexcoord = self.model.mesh_facetexcoord.flatten()
mesh_facetexcoord = mesh_facetexcoord[
mesh_facetexcoord_ranges[self.dataid] : mesh_facetexcoord_ranges[
self.dataid + 1
]
]
mesh_facetexcoord[mesh_facetexcoord == len(mesh_texcoord)] = 0
return mesh_texcoord, mesh_facetexcoord
def _get_mesh_geometry(self):
mesh_vert_adr_from = self.model.mesh_vertadr[self.dataid]
mesh_vert_adr_to = (
self.model.mesh_vertadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_vert)
)
mesh_vert = self.model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
mesh_face_adr_from = self.model.mesh_faceadr[self.dataid]
mesh_face_adr_to = (
self.model.mesh_faceadr[self.dataid + 1]
if self.dataid < self.model.nmesh - 1
else len(self.model.mesh_face)
)
mesh_face = self.model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
mesh_facenum = self.model.mesh_facenum[self.dataid]
return mesh_vert, mesh_face, mesh_facenum
class USDPrimitiveMesh(USDObject):
"""Class to handle primitive shapes in the USD scene."""
def __init__(
self,
mesh_config: dict[any, any],
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
self.mesh_config = mesh_config
self.prim_mesh = self.generate_primitive_mesh()
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)
mesh_vert, mesh_face, mesh_facenum = self._get_mesh_geometry()
self.usd_mesh.GetPointsAttr().Set(mesh_vert)
self.usd_mesh.GetFaceVertexCountsAttr().Set(
[3 for _ in range(mesh_facenum)]
)
self.usd_mesh.GetFaceVertexIndicesAttr().Set(mesh_face)
# setting mesh uv properties
mesh_texcoord, _ = 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(list(range(mesh_facenum * 3))))
self._set_refinement_properties(self.usd_prim)
if self.texture_file:
self.attach_image_material(self.usd_mesh)
else:
self.attach_solid_material(self.usd_mesh)
def generate_primitive_mesh(self):
"""Generates the mesh for the primitive USD object."""
_, prim_mesh = shapes_component.mesh_generator(self.mesh_config)
prim_mesh.translate(-prim_mesh.get_center())
return prim_mesh
def _get_uv_geometry(self):
assert self.prim_mesh
mesh_texcoord = np.array(self.prim_mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(self.prim_mesh.triangles)
# TODO(etom): bring back support for rescaling the texture coordinates.
return mesh_texcoord, mesh_facetexcoord.flatten()
def _get_mesh_geometry(self):
assert self.prim_mesh
# get mesh geometry from the open3d mesh model
mesh_vert = np.asarray(self.prim_mesh.vertices)
mesh_face = np.asarray(self.prim_mesh.triangles)
return mesh_vert, mesh_face, len(mesh_face)
class USDTendon(USDObject):
"""Class to handle tendons in the USD scene."""
def __init__(
self,
mesh_config: dict[any, any],
stage: Usd.Stage,
geom: mujoco.MjvGeom,
obj_name: str,
rgba: np.ndarray = np.array([1, 1, 1, 1]),
texture_file: Optional[str] = None,
):
super().__init__(stage, geom, obj_name, rgba, texture_file)
self.mesh_config = mesh_config
self.tendon_parts = self.generate_primitive_mesh()
self.usd_refs = collections.defaultdict(dict)
for name, _ in self.tendon_parts.items():
part_xform_path = f"{self.xform_path}/Mesh_Xform_{name}"
mesh_path = f"{part_xform_path}/Mesh_{obj_name}"
usd_xform = UsdGeom.Xform.Define(stage, part_xform_path)
self.usd_refs[name]["usd_xform"] = usd_xform
self.usd_refs[name]["usd_mesh"] = UsdGeom.Mesh.Define(stage, mesh_path)
self.usd_refs[name]["usd_prim"] = stage.GetPrimAtPath(mesh_path)
# adding ops for each of the part xforms
self.usd_refs[name]["translate_op"] = usd_xform.AddTranslateOp()
self.usd_refs[name]["scale_op"] = usd_xform.AddScaleOp()
# setting mesh geometry properties for each of the parts in the tendon
part_geometries = self._get_mesh_geometry()
for name, part_geometry in part_geometries.items():
self.usd_refs[name]["usd_mesh"].GetPointsAttr().Set(
part_geometry["mesh_vert"]
)
self.usd_refs[name]["usd_mesh"].GetFaceVertexCountsAttr().Set(
[3 for _ in range(part_geometry["mesh_facenum"])]
)
self.usd_refs[name]["usd_mesh"].GetFaceVertexIndicesAttr().Set(
part_geometry["mesh_face"]
)
# setting uv properties for each of the parts in the tendon
part_uv_geometries = self._get_uv_geometry()
for name, part_uv_geometry in part_uv_geometries.items():
self.texcoords = UsdGeom.PrimvarsAPI(
self.usd_refs[name]["usd_mesh"]
).CreatePrimvar(
"UVMap",
Sdf.ValueTypeNames.TexCoord2fArray,
UsdGeom.Tokens.faceVarying,
)
self.texcoords.Set(part_uv_geometry["mesh_texcoord"])
self.texcoords.SetIndices(
Vt.IntArray(list(range(part_geometry["mesh_facenum"] * 3)))
)
for _, ref in self.usd_refs.items():
self._set_refinement_properties(ref["usd_prim"])
if self.texture_file:
self.attach_image_material(ref["usd_mesh"])
else:
self.attach_solid_material(ref["usd_mesh"])
def generate_primitive_mesh(self):
"""Generates the tendon mesh using primitives."""
mesh_parts = {}
for part_config in self.mesh_config:
mesh_name, prim_mesh = shapes_component.mesh_generator(part_config)
prim_mesh.translate(-prim_mesh.get_center())
mesh_parts[mesh_name] = prim_mesh
return mesh_parts
def _get_uv_geometry(self):
part_uv_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items():
mesh_texcoord = np.array(mesh.triangle_uvs)
mesh_facetexcoord = np.asarray(mesh.triangles)
part_uv_geometries[name] = {
"mesh_texcoord": mesh_texcoord,
"mesh_facetexcoord": mesh_facetexcoord,
}
return part_uv_geometries
def _get_mesh_geometry(self):
part_geometries = collections.defaultdict(dict)
for name, mesh in self.tendon_parts.items():
# get mesh geometry from the open3d mesh model
mesh_vert = np.asarray(mesh.vertices)
mesh_face = np.asarray(mesh.triangles)
part_geometries[name] = {
"mesh_vert": mesh_vert,
"mesh_face": mesh_face,
"mesh_facenum": len(mesh_face),
}
return part_geometries
def update(
self,
pos: np.ndarray,
mat: np.ndarray,
visible: bool,
frame: int,
scale: Optional[np.ndarray] = None,
):
"""Updates the position and orientation of an object in the scene."""
super().update(pos, mat, visible, frame, scale)
for name in self.tendon_parts:
if "left" in name:
translate = [0, 0, -scale[2] - (scale[0] / 2)]
self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
elif "right" in name:
translate = [0, 0, scale[2] + (scale[0] / 2)]
self.usd_refs[name]["translate_op"].Set(Gf.Vec3f(translate), frame)
def update_scale(self, scale: np.ndarray, frame: int):
"""Updates the scale of the tendon."""
for name in self.tendon_parts:
if "cylinder" in name:
self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(scale.tolist()), frame)
else:
hemisphere_scale = scale.tolist()
hemisphere_scale[2] = hemisphere_scale[0]
self.usd_refs[name]["scale_op"].Set(Gf.Vec3f(hemisphere_scale), frame)
python/mujoco/usd/shapes.py
+191
View File
@@ -0,0 +1,191 @@
# Copyright 2024 DeepMind Technologies Limited
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
# http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""Built-in shapes for USD exporter."""
import mujoco
import numpy as np
import open3d as o3d
def create_hemisphere(
radius: float, theta_steps: int = 50, phi_steps: int = 50
):
"""Creates a hemisphere mesh from a point cloud."""
points = []
for i in range(phi_steps + 1):
phi = np.pi / 2 * i / phi_steps
for j in range(theta_steps + 1):
theta = 2 * np.pi * j / theta_steps
x = radius * np.sin(phi) * np.cos(theta)
y = radius * np.sin(phi) * np.sin(theta)
z = radius * np.cos(phi)
points.append([x, y, z])
pcd = o3d.geometry.PointCloud()
pcd.points = o3d.utility.Vector3dVector(points)
mesh = pcd.compute_convex_hull()[0]
return mesh
def decouple_config(config: dict[str, any]):
"""Breaks a shape config into is subcomponent shapes."""
decoupled_config = []
for key, value in config.items():
if key == "name":
continue
decoupled_config.append({
"parent_name": config["name"],
"name": config["name"] + "_" + key,
key: value.copy(),
})
return decoupled_config
def mesh_config_generator(
name: str,
geom_type: mujoco.mjtGeom,
size: np.ndarray,
decouple: bool = False,
):
"""Creates a config for a particular mesh."""
if geom_type == mujoco.mjtGeom.mjGEOM_PLANE:
config = {
"name": name,
"box": {
"width": size[0] * 2 if size[0] > 0 else 100,
"height": size[1] * 2 if size[1] > 0 else 100,
"depth": 0.001,
"map_texture_to_each_face": True,
},
}
elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE:
config = {"name": name, "sphere": {"radius": float(size[0])}}
elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE:
cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size)
config = {
"name": name,
"cylinder": cylinder["cylinder"],
"left_hemisphere": {
"radius": size[0],
"transform": {
"translate": (0, 0, -size[2]),
"rotate": (np.pi, 0, 0),
},
},
"right_hemisphere": {
"radius": size[0],
"transform": {"translate": (0, 0, size[2])},
},
}
elif geom_type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
sphere = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_SPHERE, [1.0])
sphere["sphere"]["transform"] = {"scale": tuple(size)}
config = {
"name": name,
"sphere": sphere["sphere"],
}
elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER:
config = {
"name": name,
"cylinder": {
"radius": size[0],
"height": size[2] * 2,
},
}
elif geom_type == mujoco.mjtGeom.mjGEOM_BOX:
config = {
"name": name,
"box": {
"width": size[0] * 2,
"height": size[1] * 2,
"depth": size[2] * 2,
},
}
else:
raise NotImplementedError(
f"{geom_type} primitive geom type not implemented with USD integration"
)
if decouple:
config = decouple_config(config)
return config
def mesh_generator(
mesh_config: dict[str, any],
resolution: int = 100,
):
"""Generates a mesh given a config consisting of shapes."""
assert "name" in mesh_config
prim_mesh, mesh = None, None
for shape, config in mesh_config.items():
if "name" in shape:
continue
if "box" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_box(
width=mesh_config[shape]["width"],
height=mesh_config[shape]["height"],
depth=mesh_config[shape]["depth"],
create_uv_map=True,
map_texture_to_each_face=True,
)
elif "hemisphere" in shape:
prim_mesh = create_hemisphere(radius=mesh_config[shape]["radius"])
elif "sphere" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_sphere(
radius=mesh_config[shape]["radius"],
resolution=resolution,
create_uv_map=True,
)
elif "cylinder" in shape:
prim_mesh = o3d.geometry.TriangleMesh.create_cylinder(
radius=mesh_config[shape]["radius"],
height=mesh_config[shape]["height"],
resolution=resolution,
create_uv_map=True,
)
if "transform" in config:
if "rotate" in config["transform"]:
rotation = np.zeros(9)
quat = np.zeros(4)
euler = config["transform"]["rotate"]
seq = "xyz"
mujoco.mju_euler2Quat(quat, euler, seq)
mujoco.mju_quat2Mat(rotation, quat)
rotation = rotation.reshape((3, 3))
prim_mesh.rotate(rotation, center=(0, 0, 0))
if "scale" in config["transform"]:
prim_mesh.vertices = o3d.utility.Vector3dVector(
np.asarray(prim_mesh.vertices)
* np.array(config["transform"]["scale"])
)
if "translate" in config["transform"]:
prim_mesh.translate(config["transform"]["translate"])
if not mesh:
mesh = prim_mesh
else:
mesh += prim_mesh
return mesh_config["name"], mesh