From ba842fc0d1d963466c7f99ae2ad4d83fa7cd9561 Mon Sep 17 00:00:00 2001 From: Abhishek Joshi Date: Sun, 2 Jun 2024 21:16:58 -0500 Subject: [PATCH] decoupling tendons to subcomponents and fixing tendon scaling issue --- python/mujoco/usd/__init__.py | 4 +- python/mujoco/usd/camera.py | 38 ++ python/mujoco/usd/exporter.py | 32 +- python/mujoco/usd/lights.py | 59 +++ python/mujoco/usd/objects.py | 497 ++++++++++++++++++++ python/mujoco/usd/shapes.py | 38 +- python/mujoco/usd/{component.py => temp.py} | 38 +- 7 files changed, 654 insertions(+), 52 deletions(-) create mode 100644 python/mujoco/usd/camera.py create mode 100644 python/mujoco/usd/lights.py create mode 100644 python/mujoco/usd/objects.py rename python/mujoco/usd/{component.py => temp.py} (96%) diff --git a/python/mujoco/usd/__init__.py b/python/mujoco/usd/__init__.py index b3c238f5..9bdfaccf 100644 --- a/python/mujoco/usd/__init__.py +++ b/python/mujoco/usd/__init__.py @@ -1,4 +1,6 @@ from .exporter import * -from .component import * +from .objects import * +from .lights import * +from .camera import * from .utils import * from .shapes import * \ No newline at end of file diff --git a/python/mujoco/usd/camera.py b/python/mujoco/usd/camera.py new file mode 100644 index 00000000..64c0103c --- /dev/null +++ b/python/mujoco/usd/camera.py @@ -0,0 +1,38 @@ +from typing import List, Optional, Tuple + +import utils as utils_component + +import numpy as np + +from pxr import Gf +from pxr import Usd +from pxr import UsdGeom + +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 = utils_component.create_transform_matrix( + rotation_matrix=cam_mat, translation_vector=cam_pos + ).T + self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame) \ No newline at end of file diff --git a/python/mujoco/usd/exporter.py b/python/mujoco/usd/exporter.py index f536de39..371ca5bf 100644 --- a/python/mujoco/usd/exporter.py +++ b/python/mujoco/usd/exporter.py @@ -17,9 +17,11 @@ import os import mujoco # import mujoco.usd.shapes as shapes_module -# import mujoco.usd.component as component_module +# import mujoco.usd.component as object_module import shapes as shapes_module -import component as component_module +import objects as object_module +import lights as light_module +import camera as camera_module import numpy as np import scipy @@ -181,7 +183,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() @@ -195,7 +196,6 @@ class USDExporter: self.updates += 1 def _load_textures(self): - # TODO: remove code once added internally to mujoco data_adr = 0 self.texture_files = [] for texture_id in tqdm.tqdm(range(self.model.ntex)): @@ -241,8 +241,9 @@ class USDExporter: texture_file = self.texture_files[geom.texid] if geom.texid != -1 else None + # 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, @@ -252,14 +253,15 @@ class USDExporter: texture_file=texture_file, ) else: + # handling tendons in our scene if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON: - # scale of a tendon is equivalent to its size, assume a unit length geom size mesh_config = shapes_module.mesh_config_generator( name=geom_name, geom_type=geom.type, - size=np.array([1.0, 1.0, 1.0]) + size=np.array([1.0, 1.0, 1.0]), + decouple=True ) - usd_geom = component_module.USDTendon( + usd_geom = object_module.USDTendon( mesh_config=mesh_config, stage=self.stage, geom=geom, @@ -267,13 +269,14 @@ class USDExporter: 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 = component_module.USDPrimitiveMesh( + usd_geom = object_module.USDPrimitiveMesh( mesh_config=mesh_config, stage=self.stage, geom=geom, @@ -298,7 +301,6 @@ class USDExporter: if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON: tendon_scale = geom.size - tendon_scale[2] /= 2 self.geom_refs[geom_name].update( pos=geom.pos, mat=geom.mat, @@ -321,7 +323,7 @@ class USDExporter: 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))) + (light_module.USDSphereLight(stage=self.stage, obj_name=str(i))) else: self.usd_lights.append(None) @@ -347,7 +349,7 @@ class USDExporter: 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, @@ -388,11 +390,11 @@ class USDExporter: ): 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=str(objid), radius=radius) new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0) elif light_type == "dome": - new_light = component_module.USDDomeLight( + new_light = light_module.USDDomeLight( stage=self.stage, obj_name=str(objid)) new_light.update(intensity=intensity, color=color, frame=0) @@ -403,7 +405,7 @@ class USDExporter: rotation_xyz: List[float], obj_name: Optional[str] = "camera_1", ): - new_camera = component_module.USDCamera( + new_camera = camera_module.USDCamera( stage=self.stage, obj_name=str(objid)) r = scipy.spatial.transform.Rotation.from_euler( diff --git a/python/mujoco/usd/lights.py b/python/mujoco/usd/lights.py new file mode 100644 index 00000000..eeee9e18 --- /dev/null +++ b/python/mujoco/usd/lights.py @@ -0,0 +1,59 @@ +from typing import List, Optional, Tuple + +import numpy as np + +from pxr import Gf +from pxr import Usd +from pxr import UsdGeom +from pxr import UsdLux + +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())) + diff --git a/python/mujoco/usd/objects.py b/python/mujoco/usd/objects.py new file mode 100644 index 00000000..af710351 --- /dev/null +++ b/python/mujoco/usd/objects.py @@ -0,0 +1,497 @@ +# 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 abc import ABC, abstractmethod +from collections import defaultdict +from typing import List, Optional, Tuple +import pprint + +import mujoco + +# import mujoco.usd.utils as utils_component +# import mujoco.usd.shapes as shapes_component +import utils as utils_component +import shapes as shapes_component + +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 USDObject(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 # not an arbitary value, forces difference greater than 1 for visibility on 0th frame + + @abstractmethod + def _get_uv_geometry(self): + """Gets UV information for an object in the scene""" + raise NotImplementedError + + @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 in the scene for a given frame""" + 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): + + 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): + + def __init__( + self, + mesh_config: dict, + 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, 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(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): + _, 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 + + x_scale, y_scale = self.geom.texrepeat + mesh_texcoord = np.array(self.prim_mesh.triangle_uvs) + mesh_facetexcoord = np.asarray(self.prim_mesh.triangles) + + x_multiplier, y_multiplier = 1, 1 + if self.geom.texuniform: + if "box" in self.mesh_config: + x_multiplier, y_multiplier = self.mesh_config["box"]["width"], self.mesh_config["box"]["height"] + elif "sphere" in self.mesh_config: + x_multiplier, y_multiplier = self.mesh_config["sphere"]["radius"], self.mesh_config["sphere"]["radius"] + + mesh_texcoord[:, 0] *= x_scale * x_multiplier + mesh_texcoord[:, 1] *= y_scale * y_multiplier + + 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): + + def __init__( + self, + mesh_config: dict, + 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 = defaultdict(dict) + + for name, mesh in self.tendon_parts.items(): + part_xform_path = f"{self.xform_path}/Mesh_Xform_{name}" + mesh_path = f"{part_xform_path}/Mesh_{obj_name}" + self.usd_refs[name]["usd_xform"] = UsdGeom.Xform.Define(stage, part_xform_path) + 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"] = self.usd_refs[name]["usd_xform"].AddTranslateOp() + self.usd_refs[name]["scale_op"] = self.usd_refs[name]["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([i for i in range(part_geometry["mesh_facenum"] * 3)])) + + for name in self.usd_refs.keys(): + self._set_refinement_properties(self.usd_refs[name]["usd_prim"]) + if self.texture_file: + self.attach_image_material(self.usd_refs[name]["usd_mesh"]) + else: + self.attach_solid_material(self.usd_refs[name]["usd_mesh"]) + + def generate_primitive_mesh(self): + 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 = defaultdict(dict) + for name, mesh in self.tendon_parts.items(): + x_scale, y_scale = self.geom.texrepeat + 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 = 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 for a given frame""" + super().update(pos, mat, visible, frame, scale) + for name in self.tendon_parts.keys(): + 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.keys(): + 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 index 66d9644b..d89fd43d 100644 --- a/python/mujoco/usd/shapes.py +++ b/python/mujoco/usd/shapes.py @@ -27,14 +27,28 @@ def create_hemisphere( return mesh +def decouple_config(config: dict): + 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 + size: np.ndarray, + decouple: bool = False ): if geom_type == mujoco.mjtGeom.mjGEOM_PLANE: - return { + config = { "name": name, "box": { "width": size[0] * 2 if size[0] > 0 else 100, @@ -44,7 +58,7 @@ def mesh_config_generator( } } elif geom_type == mujoco.mjtGeom.mjGEOM_SPHERE: - return { + config = { "name": name, "sphere": { "radius": float(size[0]) @@ -52,7 +66,7 @@ def mesh_config_generator( } elif geom_type == mujoco.mjtGeom.mjGEOM_CAPSULE: cylinder = mesh_config_generator(name, mujoco.mjtGeom.mjGEOM_CYLINDER, size) - return { + config = { "name": name, "cylinder": cylinder["cylinder"], "left_hemisphere": { @@ -74,12 +88,12 @@ def mesh_config_generator( sphere["sphere"]["transform"] = { "scale": tuple(size) } - return { + config = { "name": name, "sphere": sphere["sphere"], } elif geom_type == mujoco.mjtGeom.mjGEOM_CYLINDER: - return { + config = { "name": name, "cylinder": { "radius": size[0], @@ -87,7 +101,7 @@ def mesh_config_generator( } } elif geom_type == mujoco.mjtGeom.mjGEOM_BOX: - return { + config = { "name": name, "box": { "width": size[0] * 2, @@ -98,6 +112,10 @@ def mesh_config_generator( else: raise NotImplemented(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, @@ -110,7 +128,7 @@ def mesh_generator( for shape, config in mesh_config.items(): - if shape == "name": + if "name" in shape: continue if "box" in shape: @@ -140,9 +158,9 @@ def mesh_generator( ) if "transform" in config: - if "rotate" in config["transform"]: - R = mesh.get_rotation_matrix_from_xyz(config["transform"]["rotate"]) + tmp = o3d.geometry.TriangleMesh.create_coordinate_frame() # TODO: add rotation matrix code gen to utils.py + R = tmp.get_rotation_matrix_from_xyz(config["transform"]["rotate"]) prim_mesh.rotate(R, center=(0, 0, 0)) if "scale" in config["transform"]: prim_mesh.vertices = o3d.utility.Vector3dVector( diff --git a/python/mujoco/usd/component.py b/python/mujoco/usd/temp.py similarity index 96% rename from python/mujoco/usd/component.py rename to python/mujoco/usd/temp.py index 28fa3efd..2daee76b 100644 --- a/python/mujoco/usd/component.py +++ b/python/mujoco/usd/temp.py @@ -12,6 +12,7 @@ # See the License for the specific language governing permissions and # limitations under the License. # ============================================================================== +from abs import ABC from typing import List, Optional, Tuple import mujoco @@ -34,7 +35,6 @@ from pxr import UsdLux from pxr import UsdShade from pxr import Vt - class USDMesh: def __init__( @@ -265,7 +265,7 @@ class USDPrimitiveMesh: self.prim_mesh = None self.transform_op = Gf.Matrix4d(1.) - _, self.prim_mesh = shapes_component.mesh_generator(mesh_config) + self.prim_mesh = self.generate_primitive_mesh() xform_path = f"/World/{self.obj_name}_Xform" mesh_path = f"{xform_path}/{obj_name}" @@ -273,8 +273,6 @@ class USDPrimitiveMesh: self.usd_mesh = UsdGeom.Mesh.Define(stage, mesh_path) self.usd_prim = stage.GetPrimAtPath(mesh_path) - 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( @@ -302,6 +300,11 @@ class USDPrimitiveMesh: self.last_visible_frame = -2 + def generate_primitive_mesh(self): + _, prim_mesh = shapes_component.mesh_generator(self.mesh_config) + prim_mesh.translate(-prim_mesh.get_center()) + return prim_mesh + def _set_refinement_properties(self): self.usd_prim.GetAttribute("subdivisionScheme").Set("none") @@ -417,7 +420,7 @@ class USDPrimitiveMesh: 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): + def update(self, pos: np.ndarray, mat: np.ndarray, visible: bool, frame: int, ): transformation_mat = utils_component.create_transform_matrix( rotation_matrix=mat, translation_vector=pos ).T @@ -438,29 +441,12 @@ class USDPrimitiveMesh: def update_scale(self, scale: np.ndarray, frame: int): self.scale_op.Set(Gf.Vec3f(scale.tolist()), frame) -class USDTendon(USDPrimitiveMesh): - +class USDTendon: + def __init__( - self, - mesh_config: List[dict], - stage: Usd.Stage, - geom: mujoco.MjvGeom, - obj_name: str, - rgba: np.ndarray = np.array([1, 1, 1, 1]), - texture_file: Optional[str] = None, + self ): - super().__init__( - mesh_config=mesh_config, - stage=stage, - geom=geom, - obj_name=obj_name, - rgba=rgba, - texture_file=texture_file - ) - - def update(self, pos: np.ndarray, mat: np.ndarray, scale: np.ndarray, visible: bool, frame: int): - super().update(pos, mat, visible, frame) - super().update_scale(scale, frame) + pass class USDSphereLight: