diff --git a/python/mujoco/testdata/usd_golden.usda b/python/mujoco/testdata/usd_golden.usda new file mode 100644 index 00000000..1a154486 --- /dev/null +++ b/python/mujoco/testdata/usd_golden.usda @@ -0,0 +1,71 @@ +#usda 1.0 +( + defaultPrim = "World" + doc = """Generated from Composed Stage of root layer +""" + endTimeCode = 1 + startTimeCode = 0 + timeCodesPerSecond = 24 + upAxis = "Z" +) + +def Xform "World" +{ + def Xform "Light_Xform_0" + { + double3 xformOp:translate + uniform token[] xformOpOrder = ["xformOp:translate"] + + def SphereLight "Light_0" + { + bool inputs:normalize = 1 + float inputs:radius = 0.3 + bool treatAsPoint = 0 + } + } + + def Xform "CubeMesh_Xform_white_box" + { + matrix4d xformOp:transform.timeSamples = { + 0: ( (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 0), (0, 0, 0, 1) ), + } + uniform token[] xformOpOrder = ["xformOp:transform"] + + def Mesh "CubeMesh_white_box" ( + apiSchemas = ["MaterialBindingAPI"] + ) + { + int[] faceVertexCounts = [3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3, 3] + int[] faceVertexIndices = [4, 7, 5, 4, 6, 7, 0, 2, 4, 2, 6, 4, 0, 1, 2, 1, 3, 2, 1, 5, 7, 1, 7, 3, 2, 3, 7, 2, 7, 6, 0, 4, 1, 1, 4, 5] + rel material:binding = + point3f[] points = [(-1, -1, -1), (1, -1, -1), (-1, -1, 1), (1, -1, 1), (-1, 1, -1), (1, 1, -1), (-1, 1, 1), (1, 1, 1)] + texCoord2f[] primvars:UVMap = [] ( + interpolation = "faceVarying" + ) + int[] primvars:UVMap:indices = [0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16, 17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33, 34, 35] + uniform token subdivisionScheme = "none" + token visibility.timeSamples = { + 0: "inherited", + } + } + } + + def "_materials" + { + def Material "Material_white_box" + { + token outputs:surface.connect = + + def Shader "Principled_BSDF" + { + uniform token info:id = "UsdPreviewSurface" + color3f inputs:diffuseColor = (1, 1, 1) + float inputs:metallic = 0.5 + float inputs:opacity = 1 + float inputs:roughness = 0.5 + token outputs:surface + } + } + } +} + diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/component.py index 02707a13..c8a80f6f 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/component.py @@ -15,6 +15,7 @@ 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. @@ -36,9 +37,9 @@ class USDMesh: stage: Usd.Stage, model: mujoco.MjModel, geom: mujoco.MjvGeom, - objid: int, + objid: str, dataid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage @@ -216,7 +217,7 @@ class USDMesh: 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( + 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) @@ -235,8 +236,8 @@ class USDPrimitiveMesh: self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage @@ -246,9 +247,9 @@ class USDPrimitiveMesh: self.texture_file = texture_file self.usd_prim = Usd.Prim() - self.usd_mesh = Usd.Mesh() - self.prim_mesh = Usd.PrimMesh() - self.transform_op = Usd.TransformOp() + 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") @@ -364,7 +365,7 @@ class USDPrimitiveMesh: UsdShade.MaterialBindingAPI(self.usd_mesh).Bind(mtl) def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int): - transformation_mat = mujoco.usd_util.create_transform_matrix( + 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) @@ -383,8 +384,8 @@ class USDPrimitive: self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): self.stage = stage @@ -395,7 +396,7 @@ class USDPrimitive: self.usd_prim = Usd.Prim() self.usd_primitive_shape = Usd.PrimitiveShape() - self.transform_op = Usd.TransformOp() + self.transform_op = Gf.Matrix4d(1.) def _set_refinement_properties(self): self.usd_prim.GetAttribute("subdivisionScheme").Set("none") @@ -481,7 +482,7 @@ class USDPrimitive: 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( + 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) @@ -500,8 +501,8 @@ class USDCapsule(USDPrimitive): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -536,8 +537,8 @@ class USDEllipsoid(USDPrimitive): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -569,8 +570,8 @@ class USDCubeMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -586,8 +587,6 @@ class USDCubeMesh(USDPrimitiveMesh): width=self.geom.size[0] * 2, height=self.geom.size[1] * 2, depth=self.geom.size[2] * 2, - create_uv_map=True, - map_texture_to_each_face=True, ) self.prim_mesh.translate(-self.prim_mesh.get_center()) @@ -623,8 +622,8 @@ class USDSphereMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -637,7 +636,7 @@ class USDSphereMesh(USDPrimitiveMesh): 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 + radius=float(self.geom.size[0]) ) self.prim_mesh.translate(-self.prim_mesh.get_center()) @@ -673,8 +672,8 @@ class USDCylinderMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -689,7 +688,6 @@ class USDCylinderMesh(USDPrimitiveMesh): 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()) @@ -725,8 +723,8 @@ class USDPlaneMesh(USDPrimitiveMesh): self, stage: Usd.Stage, geom: mujoco.MjvGeom, - objid: int, - rgba: Tuple[int, ...] = (1, 1, 1, 1), + objid: str, + rgba: np.ndarray = np.array([1, 1, 1, 1]), texture_file: Optional[str] = None, ): @@ -742,8 +740,6 @@ class USDPlaneMesh(USDPrimitiveMesh): 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()) @@ -776,7 +772,7 @@ class USDPlaneMesh(USDPrimitiveMesh): class USDSphereLight: def __init__( - self, stage: Usd.Stage, objid: int, radius: Optional[float] = 0.3 + self, stage: Usd.Stage, objid: str, radius: Optional[float] = 0.3 ): self.stage = stage @@ -808,7 +804,7 @@ class USDSphereLight: class USDDomeLight: - def __init__(self, stage: Usd.Stage, objid: int): + def __init__(self, stage: Usd.Stage, objid: str): self.stage = stage xform_path = f"/World/Light_Xform_{objid}" @@ -829,7 +825,7 @@ class USDDomeLight: class USDCamera: - def __init__(self, stage: Usd.Stage, objid: int): + def __init__(self, stage: Usd.Stage, objid: str): self.stage = stage xform_path = f"/World/Camera_Xform_{objid}" @@ -851,7 +847,7 @@ class USDCamera: def update(self, cam_pos: np.ndarray, cam_mat: np.ndarray, frame: int): - transformation_mat = mujoco.usd_util.create_transform_matrix( + transformation_mat = mujoco.usd.utils.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/exporter.py b/python/mujoco/usd/exporter.py index c45979ef..4ec94bf6 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -15,6 +15,7 @@ import os import mujoco +import mujoco.usd.component as component_module import numpy as np import scipy import termcolor @@ -235,7 +236,7 @@ class USDExporter: # handles meshes in scene if geom.type == mujoco.mjtGeom.mjGEOM_MESH: - usd_geom = mujoco.usd.USDMesh( + usd_geom = component_module.USDMesh( stage=self.stage, model=self.model, geom=geom, @@ -245,7 +246,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE: - usd_geom = mujoco.usd.USDPlaneMesh( + usd_geom = component_module.USDPlaneMesh( stage=self.stage, geom=geom, objid=geom_name, @@ -253,7 +254,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE: - usd_geom = mujoco.usd.USDSphereMesh( + usd_geom = component_module.USDSphereMesh( stage=self.stage, geom=geom, objid=geom_name, @@ -261,7 +262,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE: - usd_geom = mujoco.usd.USDCapsule( + usd_geom = component_module.USDCapsule( stage=self.stage, geom=geom, objid=geom_name, @@ -269,7 +270,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID: - usd_geom = mujoco.usd.USDEllipsoid( + usd_geom = component_module.USDEllipsoid( stage=self.stage, geom=geom, objid=geom_name, @@ -277,7 +278,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER: - usd_geom = mujoco.usd.USDCylinderMesh( + usd_geom = component_module.USDCylinderMesh( stage=self.stage, geom=geom, objid=geom_name, @@ -285,7 +286,7 @@ class USDExporter: texture_file=texture_file, ) elif geom.type == mujoco.mjtGeom.mjGEOM_BOX: - usd_geom = mujoco.usd.USDCubeMesh( + usd_geom = component_module.USDCubeMesh( stage=self.stage, geom=geom, objid=geom_name, @@ -331,7 +332,8 @@ class USDExporter: 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(mujoco.usd.USDSphereLight(stage=self.stage, objid=i)) + self.usd_lights.append + (component_module.USDSphereLight(stage=self.stage, objid=str(i))) else: self.usd_lights.append(None) @@ -342,7 +344,7 @@ class USDExporter: if np.allclose(light.pos, [0, 0, 0]): continue - if self.usd_lights[i] is None: + if i >= len(self.usd_lights) or self.usd_lights[i] is None: continue self.usd_lights[i].update( @@ -356,8 +358,10 @@ class USDExporter: def _load_cameras(self): self.usd_cameras = [] - for name in self.camera_names: - self.usd_cameras.append(mujoco.usd.USDCamera(stage=self.stage, objid=name)) + if self.camera_names is not None: + for name in self.camera_names: + self.usd_cameras.append( + component_module.USDCamera(stage=self.stage, objid=name)) def _update_cameras( self, @@ -397,11 +401,11 @@ class USDExporter: ): if light_type == "sphere": - new_light = mujoco.usd.USDSphereLight(stage=self.stage, objid=objid, radius=radius) + new_light = component_module.USDSphereLight(stage=self.stage, objid=str(objid), radius=radius) - new_light.update(pos=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 = mujoco.usd.USDDomeLight(stage=self.stage, objid=objid) + new_light = component_module.USDDomeLight(stage=self.stage, objid=str(objid)) new_light.update(intensity=intensity, color=color, frame=0) @@ -411,11 +415,11 @@ class USDExporter: rotation_xyz: List[float], objid: Optional[int] = 1, ): - new_camera = mujoco.usd.USDCamera(stage=self.stage, objid=objid) + new_camera = component_module.USDCamera(stage=self.stage, objid=str(objid)) r = scipy.spatial.transform.Rotation.from_euler( "xyz", rotation_xyz, degrees=True) - new_camera.update(cam_pos=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"): self.stage.SetEndTimeCode(self.frame_count) diff --git a/python/mujoco/usd/exporter_test.py b/python/mujoco/usd/exporter_test.py new file mode 100644 index 00000000..17d6e855 --- /dev/null +++ b/python/mujoco/usd/exporter_test.py @@ -0,0 +1,70 @@ +# 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. +# ============================================================================== +"""Tests for the MuJoCo USD Exporter.""" + +import logging +import os + +import mujoco + +from absl.testing import absltest +from etils import epath + +# 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 +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): + + def test_usd_export(self): + if not execute_test: + return + + output_dir = os.getenv('TEST_UNDECLARED_OUTPUTS_DIR') + xml = """ + + + + + + +""" + model = mujoco.MjModel.from_xml_string(xml) + data = mujoco.MjData(model) + exporter = exporter_module.USDExporter( + model, + output_directory_name="mujoco_usdpkg", + output_directory_root=output_dir, + ) + exporter.update_scene(data) + exporter.save_scene("export.usda") + + with open(os.path.join( + 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: + self.assertEqual(f.read(), golden_file.read()) + + +if __name__ == "__main__": + absltest.main()