diff --git a/.gitignore b/.gitignore index c55f0c71..c924037f 100644 --- a/.gitignore +++ b/.gitignore @@ -18,6 +18,7 @@ *.egg-info/ build/ build_cmake/ +python/dist/ # Exclude macOS folder attributes .DS_Store diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py new file mode 100644 index 00000000..71b8c4fe --- /dev/null +++ b/python/mujoco/usd/camera.py @@ -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) diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py deleted file mode 100644 index 77f3e21e..00000000 --- a/python/mujoco/usd/component.py +++ /dev/null @@ -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) diff --git a/python/mujoco/usd/demo.py b/python/mujoco/usd/demo.py new file mode 100644 index 00000000..442e33e3 --- /dev/null +++ b/python/mujoco/usd/demo.py @@ -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) diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index 8ad10af8..cdd0f9c8 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -12,27 +12,30 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +"""USD exporter.""" + import os +from typing import List, Optional 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 +from PIL import Image as im +from PIL import ImageOps import scipy import termcolor 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 Usd from pxr import UsdGeom class USDExporter: + """MuJoCo to USD exporter for porting scenes to external renderers.""" def __init__( self, @@ -47,13 +50,13 @@ class USDExporter: specialized_materials_file: Optional[str] = None, verbose: bool = True, ): - """Initializes a new USD Exporter + """Initializes a new USD Exporter. Args: model: an MjModel instance. height: image height 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 automatically based on the estimated maximum number of renderable geoms in the model. @@ -61,6 +64,10 @@ class USDExporter: and assets generated by the USD renderer. output_directory_root: path to root directory storing generated frames 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. """ @@ -99,7 +106,12 @@ class USDExporter: self.frame_count = 0 # maintains how many times we have saved the scene 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 self.renderer = mujoco.Renderer(model, height, width, max_geom) @@ -114,18 +126,21 @@ class USDExporter: @property def usd(self): + """Returns the USD file as a string.""" return self.stage.GetRootLayer().ExportToString() @property def scene(self): + """Returns the scene.""" return self.renderer.scene def _initialize_usd_stage(self): + """Initializes a USD stage to represent the mujoco scene.""" self.stage = Usd.Stage.CreateInMemory() UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z) self.stage.SetStartTimeCode(0) # add as user input - self.stage.SetTimeCodesPerSecond(24.0) + self.stage.SetTimeCodesPerSecond(60.0) default_prim = UsdGeom.Xform.Define( self.stage, Sdf.Path("/World") @@ -133,6 +148,7 @@ class USDExporter: self.stage.SetDefaultPrim(default_prim) def _initialize_output_directories(self): + """Initializes output directories to store frames and assets.""" self.output_directory_path = os.path.join( self.output_directory_root, self.output_directory_name ) @@ -161,7 +177,7 @@ class USDExporter: data: mujoco.MjData, scene_option: Optional[mujoco.MjvOption] = None, ): - """Updates the scene with latest sim data + """Updates the scene with latest sim data. Args: data: structure storing current simulation state @@ -175,7 +191,6 @@ class USDExporter: # update the mujoco renderer self.renderer.update_scene(data, scene_option=scene_option) - # TODO: update scene options if self.updates == 0: self._initialize_usd_stage() @@ -189,7 +204,7 @@ class USDExporter: self.updates += 1 def _load_textures(self): - # TODO: remove code once added internally to mujoco + """Load textures.""" data_adr = 0 self.texture_files = [] for texture_id in tqdm.tqdm(range(self.model.ntex)): @@ -212,7 +227,7 @@ class USDExporter: ) img_path = os.path.join( relative_path, texture_file_name - ) # relative path, TODO: switch back to this + ) self.texture_files.append(img_path) @@ -228,19 +243,20 @@ class USDExporter: ) 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_name2usd + assert geom_name not in self.geom_names 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 else None ) - # handles meshes in scene + # handling meshes in our scene if geom.type == mujoco.mjtGeom.mjGEOM_MESH: - usd_geom = component_module.USDMesh( + usd_geom = object_module.USDMesh( stage=self.stage, model=self.model, geom=geom, @@ -249,95 +265,77 @@ class USDExporter: rgba=geom.rgba, 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: - 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): - - geom_names = set(self.geom_name2usd.keys()) - - # iterate through all geoms in the scene and makes update + """Iterate through all geoms in the scene and makes update.""" for i in range(self.scene.ngeom): 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) - if self.geom_name2usd[geom_name]: - self.geom_name2usd[geom_name].update_visibility(False, 0) - if self.geom_name2usd[geom_name]: - self.geom_name2usd[geom_name].update( + if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON: + 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, mat=geom.mat, visible=geom.rgba[3] > 0, 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): # initializes an usd light object for every light in the scene - self.usd_lights = [] for i in range(self.scene.nlight): light = self.scene.lights[i] if not np.allclose(light.pos, [0, 0, 0]): - self.usd_lights.append - (component_module.USDSphereLight(stage=self.stage, obj_name=str(i))) + self.usd_lights.append( + light_module.USDSphereLight(stage=self.stage, obj_name=str(i)) + ) else: self.usd_lights.append(None) @@ -359,19 +357,23 @@ class USDExporter: ) def _load_cameras(self): - self.usd_cameras = [] if self.camera_names is not None: for name in self.camera_names: 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( self, data: mujoco.MjData, 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_name = self.camera_names[i] @@ -386,12 +388,14 @@ class USDExporter: up = avg_camera.up right = np.cross(forward, up) - R = np.eye(3) - R[:, 0] = right - R[:, 1] = up - R[:, 2] = -forward + rotation = np.eye(3) + rotation[:, 0] = right + rotation[:, 1] = up + 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( self, @@ -399,37 +403,91 @@ class USDExporter: intensity: int, radius: Optional[float] = 1.0, 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", ): + """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": - 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": - new_light = component_module.USDDomeLight( - stage=self.stage, obj_name=str(objid)) - - new_light.update(intensity=intensity, color=color, frame=0) + new_light = light_module.USDDomeLight(stage=self.stage, obj_name=obj_name) + new_light.update(intensity=intensity, color=color) def add_camera( self, pos: List[float], rotation_xyz: List[float], - objid: Optional[int] = 1, + obj_name: Optional[str] = "camera_1", ): - new_camera = component_module.USDCamera( - stage=self.stage, obj_name=str(objid)) + """Adds a user defined, fixed camera. + + 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( "xyz", rotation_xyz, degrees=True) new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0) 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) + + # 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( - 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: - 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) diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py index 17d6e855..01424bbf 100644 --- a/python/mujoco/usd/exporter_test.py +++ b/python/mujoco/usd/exporter_test.py @@ -17,21 +17,19 @@ import logging import os -import mujoco - from absl.testing import absltest from etils import epath +import mujoco + # Open3D and USD are not fully supported on all MuJoCo architectures. -# pylint: disable=python.style(g-import-not-at-top) execute_test = True try: - from mujoco.usd import exporter as exporter_module - from pxr import Usd + from mujoco.usd import exporter as exporter_module # pylint: disable=g-import-not-at-top except ImportError: logging.warning('Skipping test due to missing import') execute_test = False -# pylint: enable=python.style(g-import-not-at-top) + class ExporterTest(absltest.TestCase): @@ -52,19 +50,19 @@ class ExporterTest(absltest.TestCase): data = mujoco.MjData(model) exporter = exporter_module.USDExporter( model, - output_directory_name="mujoco_usdpkg", + output_directory_name='mujoco_usdpkg', output_directory_root=output_dir, ) exporter.update_scene(data) - exporter.save_scene("export.usda") + exporter.save_scene('export.usda') 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( - epath.resource_path("mujoco"), "testdata", "usd_golden.usda") - with open(golden_path, "r") as golden_file: + epath.resource_path('mujoco'), 'testdata', 'usd_golden.usda') + with open(golden_path, 'r') as golden_file: self.assertEqual(f.read(), golden_file.read()) -if __name__ == "__main__": +if __name__ == '__main__': absltest.main() diff --git a/python/mujoco/usd/lights.py b/python/mujoco/usd/lights.py new file mode 100644 index 00000000..6a0d9d41 --- /dev/null +++ b/python/mujoco/usd/lights.py @@ -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())) diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py new file mode 100644 index 00000000..579a1d46 --- /dev/null +++ b/python/mujoco/usd/objects.py @@ -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) diff --git a/python/mujoco/usd/shapes.py b/python/mujoco/usd/shapes.py new file mode 100644 index 00000000..c4308cad --- /dev/null +++ b/python/mujoco/usd/shapes.py @@ -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