diff --git a/python/dist/mujoco-3.0.2.tar.gz b/python/dist/mujoco-3.0.2.tar.gz
new file mode 100644
index 00000000..cd71ac7c
Binary files /dev/null and b/python/dist/mujoco-3.0.2.tar.gz differ
diff --git a/python/mujoco/usd_component.py b/python/mujoco/usd_component.py
index afd05133..c649096e 100644
--- a/python/mujoco/usd_component.py
+++ b/python/mujoco/usd_component.py
@@ -1,407 +1,632 @@
-import os
-from enum import Enum
+import numpy as np
+import open3d as o3d
+from usd_utils import *
+from mujoco import mjtGeom
+from pxr import Gf, Sdf, Vt
+from typing import Optional, List
+from mujoco import _structs, _enums
+from pxr import Usd, UsdGeom, UsdShade, UsdLux
-import mujoco
-from mujoco.usd_utilities import *
-from pxr import Usd, UsdGeom, UsdLux, UsdShade, Vt, Gf, Sdf
-from scipy.spatial.transform import Rotation as R
+class USDMesh:
-# TODO: clean this up and remove the if statements
-def create_usd_geom_primitive(geom, stage, texture_file):
- geom_type = geom.type
- if geom_type==USDGeomType.Plane.value:
- return USDPlane(geom, stage, texture_file)
- elif geom_type==USDGeomType.Sphere.value:
- return USDSphere(geom, stage, texture_file)
- elif geom_type==USDGeomType.Capsule.value:
- return USDCapsule(geom, stage, texture_file)
- elif geom_type==USDGeomType.Cylinder.value:
- return USDCylinder(geom, stage, texture_file)
- elif geom_type==USDGeomType.Cube.value:
- return USDCube(geom, stage, texture_file)
- elif geom_type==USDGeomType.Cube.value:
- return USDCube(geom, stage, texture_file)
- else:
- return None
+ def __init__(
+ self,
+ stage: Usd.Stage,
+ model: _structs.MjModel,
+ geom: _structs.MjvGeom,
+ objid: int,
+ dataid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None
+ ):
+ """ Initializes a new USD mesh
+ Args:
+ model: an MjModel instance.
+ dataid: id of the mesh
+ texture_file: texture associated with the mesh
+ """
+ self.stage = stage
+ self.model = model
+ self.geom = geom
+ self.objid = objid
+ self.rgba = rgba
+ self.dataid = dataid
+ self.texture_file = texture_file
-class USDGeomType(Enum):
- """
- Represents different types of geoms we can add to USD
- The values match those found by the enum presented here:
- https://mujoco.readthedocs.io/en/latest/APIreference/APItypes.html#mjtgeom
- """
- Plane = 0
- # Hfield = 1
- Sphere = 2
- Capsule = 3
- # Ellipsoid = 4
- Cylinder = 5
- Cube = 6
- Mesh = 7
+ xform_path = f'/World/Mesh_Xform_{objid}'
+ mesh_path= f'{xform_path}/Mesh_{objid}'
+ 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)
-class USDGeom(object):
- """
- Parent class for created geoms
- """
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- self.geom = geom
- self.stage = stage
- self.texture_file = texture_file
- self.type = None
- self.xform = None
- self.prim = None
- self.ref = None
+ # 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)
- # xform operations, set when updating geoms
- self.translate_op = None
- self.rotate_op = None
- self.scale_op = None
+ # 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()))
- def set_texture(self):
- if self.texture_file:
- mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(self.texture_file))[0]}")
- mtl = UsdShade.Material.Define(self.stage, mtl_path)
- shader = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("Shader"))
- shader.CreateIdAttr("UsdPreviewSurface")
- shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0))
- shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
- shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+ self._attach_material()
+
+ # defining ops required by update function
+ self.transform_op = self.usd_xform.AddTransformOp()
- diffuse_tx = UsdShade.Shader.Define(self.stage, mtl_path.AppendPath("DiffuseColorTx"))
- diffuse_tx.CreateIdAttr('UsdUVTexture')
+ 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
- diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(self.texture_file)
- diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3)
- shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb')
- mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
+ 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]
- self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
- UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
+ 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]]
- def update_geom(self, new_geom):
- self.update_pos(new_geom.pos)
- self.update_rotation(new_geom.mat)
- self.update_size(new_geom.size)
- self.update_color(new_geom.rgba)
- self.update_transparency(new_geom.rgba[3])
-
- def update_pos(self, new_pos):
- pos = tuple([float(x) for x in new_pos])
- if not self.translate_op:
- self.translate_op = self.xform.AddTranslateOp()
- self.translate_op.Set(pos)
-
- def update_rotation(self, new_mat):
- r = R.from_matrix(new_mat)
- euler_rotation = r.as_euler('xyz', degrees=True)
- rotation = Gf.Vec3f(float(euler_rotation[0]), float(euler_rotation[1]), float(euler_rotation[2]))
- if not self.rotate_op:
- self.rotate_op = self.xform.AddRotateXYZOp()
- self.rotate_op.Set(rotation)
-
- # TODO: check to make sure scale and size are the same thing
- def update_size(self, new_size):
- size = tuple([float(x) for x in new_size])
- if not self.scale_op:
- self.scale_op = self.xform.AddScaleOp()
- self.scale_op.Set(value=size)
-
- def update_color(self, new_color):
- # new_color is the rgba (we extract first three)
- rgba = [(float(x) for x in new_color[:3])]
- self.prim.GetDisplayColorAttr().Set(rgba)
-
- def update_transparency(self, new_transparency):
- self.prim.GetDisplayOpacityAttr().Set([float(abs(new_transparency))])
- if new_transparency < 0:
- self.prim.MakeInvisible()
-
- def __str__(self):
- return f'type = {self.type} \ngeom = {self.geom}'
-
-class USDPlane(USDGeom):
- """
- Stores information regarding a plane geom in USD
- """
-
- plane_count = 0
-
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- super().__init__(geom, stage, texture_file)
- self.type = 0
- USDPlane.plane_count += 1
- xform_path = f'/World/Plane_Xform_{USDPlane.plane_count}'
- plane_path = f'{xform_path}/Plane_{USDPlane.plane_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Cube.Define(stage, plane_path) # temporary fix for planes
- self.ref = stage.GetPrimAtPath(plane_path)
-
- self.set_texture()
-
- def update_size(self, new_size):
- size = tuple([float(new_size[0]), float(new_size[1]), 0.005])
- if not self.scale_op:
- self.scale_op = self.xform.AddScaleOp()
- self.scale_op.Set(value=size)
-
-class USDSphere(USDGeom):
- """
- Stores information regarding a sphere geom in USD
- """
-
- sphere_count = 0
-
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- super().__init__(geom, stage, texture_file)
- self.type = 2
- USDSphere.sphere_count += 1
- xform_path = f'/World/Sphere_Xform_{USDSphere.sphere_count}'
- sphere_path = f'{xform_path}/Sphere_{USDSphere.sphere_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Sphere.Define(stage, sphere_path)
- self.ref = stage.GetPrimAtPath(sphere_path)
-
- self.set_texture()
-
-class USDCapsule(USDGeom):
- """
- Stores information regarding a capsule geom in USD
- """
-
- capsule_count = 0
-
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- super().__init__(geom, stage, texture_file)
- self.type = 3
- USDCapsule.capsule_count += 1
- xform_path = f'/World/Capsule_Xform_{USDCapsule.capsule_count}'
- capsule_path = f'{xform_path}/Capsule_{USDCapsule.capsule_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Capsule.Define(stage, capsule_path)
- self.ref = stage.GetPrimAtPath(capsule_path)
-
- self.set_texture()
-
-class USDCylinder(USDGeom):
- """
- Stores information regarding a capsule geom in USD
- """
-
- cylinder_count = 0
-
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- super().__init__(geom, stage, texture_file)
- self.type = 5
- USDCylinder.cylinder_count += 1
- xform_path = f'/World/Cylinder_Xform_{USDCylinder.cylinder_count}'
- cylinder_path = f'{xform_path}/Cylinder_{USDCylinder.cylinder_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Cylinder.Define(stage, cylinder_path)
- self.ref = stage.GetPrimAtPath(cylinder_path)
-
- self.set_texture()
-
-class USDCube(USDGeom):
- """
- Stores information regarding a cube geom in USD
- """
-
- cube_count = 0
-
- def __init__(self,
- geom=None,
- stage=None,
- texture_file=None):
- super().__init__(geom, stage, texture_file)
- self.type = 6
- USDCube.cube_count += 1
- xform_path = f'/World/Cube_Xform_{USDCube.cube_count}'
- cube_path = f'{xform_path}/Cube_{USDCube.cube_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Cube.Define(stage, cube_path)
- self.ref = stage.GetPrimAtPath(cube_path)
-
- self.set_texture()
-
-class USDMesh(USDGeom):
- """
- Stores information regarding a mesh geom in USD
- """
-
- mesh_count = 0
-
- def __init__(self,
- mesh_idx,
- geom,
- stage,
- model,
- texture_file):
- super().__init__(geom, stage)
-
- assert mesh_idx != -1
-
- mesh_vert_adr_from = model.mesh_vertadr[mesh_idx]
- mesh_vert_adr_to = model.mesh_vertadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_vert)
- mesh_vert = model.mesh_vert[mesh_vert_adr_from:mesh_vert_adr_to]
-
- mesh_face_adr_from = model.mesh_faceadr[mesh_idx]
- mesh_face_adr_to = model.mesh_faceadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_face)
- mesh_face = model.mesh_face[mesh_face_adr_from:mesh_face_adr_to]
-
- self.type = 7
- xform_path = f'/World/Mesh_Xform_{USDMesh.mesh_count}'
- mesh_path= f'{xform_path}/Mesh_{USDMesh.mesh_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdGeom.Mesh.Define(stage, mesh_path)
- self.ref = stage.GetPrimAtPath(mesh_path)
-
- self.vertices = mesh_vert
- self.prim.GetPointsAttr().Set(self.vertices)
-
- model.mesh_facenum[mesh_idx]
- self.prim.GetFaceVertexCountsAttr().Set([3 for _ in range(model.mesh_facenum[mesh_idx])])
-
- self.faces = mesh_face
- self.prim.GetFaceVertexIndicesAttr().Set(self.faces)
-
- self.texture_file = texture_file
- if texture_file:
-
- mesh_texcoord_adr_from = model.mesh_texcoordadr[mesh_idx]
- mesh_texcoord_adr_to = model.mesh_texcoordadr[mesh_idx+1] if mesh_idx < model.nmesh - 1 else len(model.mesh_texcoord)
- mesh_texcoord = model.mesh_texcoord[mesh_texcoord_adr_from:mesh_texcoord_adr_to]
-
- # texid = geom.texid
- # texcoords = model.mesh_texcoord[mesh_texcoord_ranges[texid]:mesh_texcoord_ranges[texid+1]]
-
- mesh_facetexcoord_ranges = get_facetexcoord_ranges(model.nmesh, model.mesh_facenum)
-
- facetexcoords = model.mesh_facetexcoord.flatten()
- facetexcoords = facetexcoords[mesh_facetexcoord_ranges[mesh_idx]:mesh_facetexcoord_ranges[mesh_idx+1]]
- self.texcoords = UsdGeom.PrimvarsAPI(self.prim).CreatePrimvar("st",
- Sdf.ValueTypeNames.TexCoord2fArray,
- UsdGeom.Tokens.faceVarying)
-
- self.texcoords.Set(mesh_texcoord)
- self.texcoords.SetIndices(Vt.IntArray(facetexcoords.tolist()));
-
- mtl_path = Sdf.Path(f"/World/Looks/Material_{os.path.splitext(os.path.basename(texture_file))[0]}")
- mtl = UsdShade.Material.Define(stage, mtl_path)
- shader = UsdShade.Shader.Define(stage, mtl_path.AppendPath("Shader"))
- shader.CreateIdAttr("UsdPreviewSurface")
- shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set((1.0, 0.0, 0.0))
- shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
- shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
-
- diffuse_tx = UsdShade.Shader.Define(stage,mtl_path.AppendPath("DiffuseColorTx"))
- diffuse_tx.CreateIdAttr('UsdUVTexture')
-
- diffuse_tx.CreateInput('file', Sdf.ValueTypeNames.Asset).Set(texture_file)
- diffuse_tx.CreateOutput('rgb', Sdf.ValueTypeNames.Float3)
- shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).ConnectToSource(diffuse_tx.ConnectableAPI(), 'rgb')
- mtl.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface")
-
- self.prim.GetPrim().ApplyAPI(UsdShade.MaterialBindingAPI)
- UsdShade.MaterialBindingAPI(self.prim).Bind(mtl)
-
- USDMesh.mesh_count += 1
-
- def update_geom(self, new_geom):
- self.update_pos(new_geom.pos)
- self.update_rotation(new_geom.mat)
- if not self.texture_file:
- self.update_color(new_geom.rgba)
- self.update_transparency(new_geom.rgba[3])
-
-class USDLight(object):
- """
- Class for the created lights
- """
-
- light_count = 0
+ return mesh_texcoord, mesh_facetexcoord
+
+ def _get_mesh_geometry(self):
+ # get mesh geometry structure from reading the mjModel
+ 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]
- def __init__(self,
- stage):
- self.stage = stage
- USDLight.light_count += 1
- xform_path = f'/World/Light_Xform_{USDLight.light_count}'
- light_path = f'{xform_path}/Light_{USDLight.light_count}'
- self.xform = UsdGeom.Xform.Define(stage, xform_path)
- self.prim = UsdLux.SphereLight.Define(stage, light_path)
- self.ref = stage.GetPrimAtPath(light_path)
-
- self.translate_op = None
-
- def update_light(self, new_light):
- pos = tuple([float(x) for x in new_light.pos])
- if not self.translate_op:
- self.translate_op = self.xform.AddTranslateOp()
- self.translate_op.Set(pos)
-
- if pos == (0, 0, 0):
- self.prim.GetIntensityAttr().Set(0);
- else:
- self.prim.GetIntensityAttr().Set(20000);
-
- # TODO attributes:
- # - direction
- # - intensity
- # - exposure
- # - radius
- # - specular
-
-class USDCamera(object):
- """
- Class for created cameras
- """
-
- camera_count = 0
-
- def __init__(self,
- stage):
- self.stage = stage
- USDCamera.camera_count += 1
- camera_path = f'/World/Camera_{USDCamera.camera_count}'
- self.prim = UsdGeom.Camera.Define(stage, camera_path)
- self.ref = stage.GetPrimAtPath(camera_path)
-
- def update_camera(self, new_pos, new_quat):
- # print("---- Updating camera in USD ----")
- xformAPI = UsdGeom.XformCommonAPI(self.prim)
-
- pos = tuple([float(x) for x in new_pos])
-
- # convert a quat to euler rotation angles
- r = R.from_quat(new_quat)
- euler_rotation = r.as_euler('xyz', degrees=True)
- rotation = Gf.Vec3f(float(euler_rotation[2]), float(euler_rotation[1]), float(euler_rotation[0]))
-
- # hardcoded values for Robosuite testing and prototype
- # TODO: use actual camera values
- xformAPI.SetTranslate(pos)
- xformAPI.SetRotate(rotation)
- # xformAPI.SetScale((1, 1, 1))
-
- self.prim.CreateFocalLengthAttr().Set(24)
- self.prim.CreateFocusDistanceAttr().Set(400)
-
-
-
-
-
-
\ No newline at end of file
+ return mesh_vert, mesh_face, mesh_facenum
+
+ def _attach_material(self):
+ mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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.array,
+ mat: np.array,
+ frame: int
+ ):
+ transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
+ self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
+
+class USDPrimitiveMesh:
+
+ def __init__(self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+ self.stage = stage
+ self.geom = geom
+ self.objid = objid
+ self.rgba = rgba
+ self.texture_file = texture_file
+
+ self.prim_mesh = None
+
+ 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)
+
+ mesh_texcoord[:, 0] *= x_scale
+ mesh_texcoord[:, 1] *= y_scale
+
+ 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.objid}")
+ mtl = UsdShade.Material.Define(self.stage, mtl_path)
+ if self.texture_file:
+ # remove all code in this if block
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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.array,
+ mat: np.array,
+ frame: int
+ ):
+ transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
+ self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
+
+class USDPrimitive:
+ def __init__(self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+ self.stage = stage
+ self.geom = geom
+ self.objid = objid
+ self.rgba = rgba
+ self.texture_file = texture_file
+
+ def _attach_material(self):
+ mtl_path = Sdf.Path(f"/World/_materials/Material_{self.objid}")
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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("roughness", Sdf.ValueTypeNames.Float).Set(0.5)
+ bsdf_shader.CreateInput("metallic", Sdf.ValueTypeNames.Float).Set(0.0)
+
+ 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.array,
+ mat: np.array,
+ frame: int
+ ):
+ transformation_mat = create_transform_matrix(rotation_matrix=mat, translation_vector=pos).T
+ self.transform_op.Set(Gf.Matrix4d(transformation_mat.tolist()), frame)
+
+class USDCapsule(USDPrimitive):
+ def __init__(self,
+ stage: Usd.Stage,
+ geom: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/Capsule_Xform_{objid}'
+ capsule_path = f'{xform_path}/Capsule_{objid}'
+ 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()
+
+ 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: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/Ellipsoid_Xform_{objid}'
+ ellipsoid_path = f'{xform_path}/Ellipsoid_{objid}'
+ 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()
+
+ 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: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/CubeMesh_Xform_{objid}'
+ mesh_path= f'{xform_path}/CubeMesh_{objid}'
+ 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,
+ 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)]))
+
+ # 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: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/SphereMesh_Xform_{objid}'
+ mesh_path= f'{xform_path}/SphereMesh_{objid}'
+ 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)]))
+
+ # 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: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/CylinderMesh_Xform_{objid}'
+ mesh_path= f'{xform_path}/CylinderMesh_{objid}'
+ 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)]))
+
+ # 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: _structs.MjvGeom,
+ objid: int,
+ rgba: List[int] = [1,1,1,1],
+ texture_file: Optional[str] = None):
+
+ super().__init__(stage,
+ geom,
+ objid,
+ rgba,
+ texture_file)
+
+ xform_path = f'/World/Plane_Xform_{objid}'
+ plane_path = f'{xform_path}/PlaneMesh_{objid}'
+ 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)]))
+
+ # setting attributes for the shape
+ self._attach_material()
+
+ # defining ops required by update function
+ self.transform_op = self.usd_xform.AddTransformOp()
+
+class USDLight:
+ def __init__(self,
+ stage: Usd.Stage,
+ objid: int,
+ radius: Optional[float] = 0.7):
+ self.stage = stage
+
+ xform_path = f'/World/Light_Xform_{objid}'
+ light_path = f'{xform_path}/Light_{objid}'
+ 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.array,
+ intensity: int,
+ color: np.array,
+ 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 USDCamera:
+ def __init__(self,
+ stage: Usd.Stage,
+ objid: int):
+ self.stage = stage
+
+ xform_path = f'/World/Camera_Xform_{objid}'
+ camera_path = f'{xform_path}/Camera_{objid}'
+ 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.CreateFocusDistanceAttr().Set(400)
+
+ self.usd_camera.GetClippingRangeAttr().Set(Gf.Vec2f(1e-4, 1e6))
+
+ def update(self,
+ cam_pos: np.array,
+ cam_mat: np.array,
+ frame: int):
+
+ transformation_mat = 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_renderer.py b/python/mujoco/usd_renderer.py
index 81be45a0..51d3733f 100644
--- a/python/mujoco/usd_renderer.py
+++ b/python/mujoco/usd_renderer.py
@@ -1,217 +1,325 @@
import os
-import shutil
-from termcolor import colored
+import pprint
import mujoco
-import mujoco.viewer as viewer
-from mujoco.usd_component import *
-from mujoco.usd_utilities import *
-from pxr import Usd, UsdGeom
-
-from mujoco import _structs
-
-from PIL import Image as im
+from usd_utils import *
from PIL import ImageOps
+from mujoco import mjtGeom
+from PIL import Image as im
+from usd_component import *
+from pxr import Usd, UsdGeom
+from termcolor import colored
+from mujoco import mjv_averageCamera
+from typing import Optional, List, Union, Tuple
+from mujoco import _structs, _constants, _enums
+from scipy.spatial.transform import Rotation as R
-class USDRenderer(object):
- """
- Renderer class that creates USD representations for mujoco scenes
- """
- def __init__(self,
- model,
- height=480,
- width=480,
- root_dir_name="usdpkg",
- root_dir_path=None,
- verbose=True):
- self.model = model
- self.root_dir_name = root_dir_name
- self.root_dir_path = root_dir_path
- self.verbose = verbose
- self.data = None
- self.renderer = mujoco.Renderer(model, height, width)
- self.reload_scene_info = True
- self.frame_count = 0
+class USDRenderer:
- self.create_output_directories()
-
- self.stage = Usd.Stage.CreateInMemory()
+ def __init__(
+ self,
+ model: _structs.MjModel,
+ height: int = 480,
+ width: int = 480,
+ max_geom: int = 10000,
+ output_directory_name: str = "mujoco_usdpkg",
+ output_directory_root: str = "./",
+ verbose: bool = True,
+ light_intensity: int = 10000
+ ):
+ """ Initializes a new USD Renderer
+ 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 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.
+ output_directory_name: name of root directory to store outputted frames and assets generated by the USD renderer.
+ output_directory_root: path to root directory storing generated frames and assets by the USD renderer.
+ verbose: decides whether to print updates.
+ """
- UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
+ buffer_width = model.vis.global_.offwidth
+ buffer_height = model.vis.global_.offheight
- geom_groups = [0,1,0,0,0,0] # Setting default geom groups for now
+ if width > buffer_width:
+ raise ValueError(f"""
+ Image width {width} > framebuffer width {buffer_width}. Either reduce the image
+ width or specify a larger offscreen framebuffer in the model XML using the
+ clause:
+
+
+ """.lstrip())
- self.scene_option = _structs.MjvOption()
- self.scene_option.geomgroup = geom_groups
+ if height > buffer_height:
+ raise ValueError(f"""
+ Image height {height} > framebuffer height {buffer_height}. Either reduce the
+ image height or specify a larger offscreen framebuffer in the model XML using
+ the clause:
+
+
+ """.lstrip())
- @property
- def usd(self):
- return self.stage.GetRootLayer().ExportToString()
-
- @property
- def scene(self):
- return self.renderer.scene
-
- def create_output_directories(self):
- if not self.root_dir_path:
- self.root_dir_path = os.getcwd()
+ self.model = model
+ self.height = height
+ self.width = width
+ self.max_geom = max_geom
+ self.output_directory_name = output_directory_name
+ self.output_directory_root = output_directory_root
+ self.verbose = verbose
+ self.light_intensity = light_intensity
- self.output_dir = os.path.join(self.root_dir_path, self.root_dir_name)
- if not os.path.exists(self.output_dir):
- os.makedirs(self.output_dir)
+ self.frame_count = 0 # maintains how many times we have saved the scene
+ self.updates = 0
- self.scenes_dir = os.path.join(self.output_dir, "scenes")
- if not os.path.exists(self.scenes_dir):
- os.makedirs(self.scenes_dir)
+ # initializing rendering requirements
+ self.renderer = mujoco.Renderer(model, height, width, max_geom)
+ self._initialize_usd_stage()
+ self._scene_option = _structs.MjvOption() # using default scene option
+
+ # initializing output_directories
+ self._initialize_output_directories()
+
+ # loading required textures for the scene
+ self._load_textures()
+
+ @property
+ def usd(self):
+ return self.stage.GetRootLayer().ExportToString()
- self.assets_dir = os.path.join(self.output_dir, "assets")
- if not os.path.exists(self.assets_dir):
- os.makedirs(self.assets_dir)
+ @property
+ def scene(self):
+ return self.renderer.scene
- if self.verbose:
- output_dir_msg = colored(f"Writing files to {self.output_dir}", "green")
- print(output_dir_msg)
+ def _initialize_usd_stage(self):
+ self.stage = Usd.Stage.CreateInMemory()
+ UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
+ self.stage.SetStartTimeCode(0)
+ # add as user imput
+ self.stage.SetTimeCodesPerSecond(60.0)
- def save_scene(self):
- output_file_path = os.path.join(self.scenes_dir, f'frame_{self.frame_count}_.usd')
- with open(output_file_path, "w") as f:
- f.write(self.usd)
- self.frame_count += 1
+ def _initialize_output_directories(self):
+ self.output_directory_path = os.path.join(self.output_directory_root, self.output_directory_name)
+ if not os.path.exists(self.output_directory_path):
+ os.makedirs(self.output_directory_path)
- def update_geom_groups(self, geom_groups):
- self.scene_option.geomgroup = geom_groups
- self.reload_scene_info = True
- self.update_scene(self.data)
+ self.frames_directory = os.path.join(self.output_directory_path, "frames")
+ if not os.path.exists(self.frames_directory):
+ os.makedirs(self.frames_directory)
+
+ self.assets_directory = os.path.join(self.output_directory_path, "assets")
+ if not os.path.exists(self.assets_directory):
+ os.makedirs(self.assets_directory)
- def update_scene(self, data):
- self.renderer.update_scene(data, scene_option=self.scene_option)
- self.data = data
+ if self.verbose:
+ print(colored(f"Writing output frames and assets to {self.output_directory_path}", "green"))
- if self.reload_scene_info:
- # loads the initial geoms, lights, and camera information
- # from the scene
- self._load()
- self.reload_scene_info = False
-
- self._update()
+ def update_scene(
+ self,
+ data: _structs.MjData,
+ camera: Union[int, str, _structs.MjvCamera] = -1,
+ scene_option: Optional[_structs.MjvOption] = None,
+ ):
+ """ Updates the scene with latest sim data
+ Args:
+ data: structure storing current simulation state
+ scene_option: we use this to determine which geom groups to activate
+ """
- def _load(self):
- """
- Loads and initializes the necessary objects to render the scene
- """
+ self.frame_count += 1
- # Create and loads the texture files to the assets directory
- # TODO: remove code once added internally to mujoco
- data_adr = 0
- texture_files = []
- for texid in range(self.model.ntex):
- height = self.model.tex_height[texid]
- width = self.model.tex_width[texid]
- pixels = 3*height*width
- rgb = self.model.tex_rgb[data_adr:data_adr+pixels]
- img = rgb.reshape(height, width, 3)
- texture_file_name = f"texture_{texid}.png"
- file_path = os.path.join(self.assets_dir, texture_file_name)
- img = im.fromarray(img)
- img = ImageOps.flip(img)
- img.save(file_path)
+ scene_option = scene_option or self._scene_option
- relative_path = os.path.relpath(self.assets_dir, self.scenes_dir)
- img_path = os.path.join(relative_path, texture_file_name)
+ # update the mujoco renderer
+ self.renderer.update_scene(data,
+ scene_option=scene_option,
+ camera=camera)
- texture_files.append(img_path)
- data_adr += pixels
+ # TODO: update scene options
+ if self.updates == 0:
+ self._initialize_usd_stage()
- # initializes an array to store all the geoms in the scene
- # populates with "empty" USDGeom objects
- self.usd_geoms = []
- geoms = self.scene.geoms
- self.ngeom = self.scene.ngeom
- for i in range(self.ngeom):
- geom = geoms[i]
- if geom.texid == -1:
- texture_file = None
- else:
- texture_file = texture_files[geom.texid]
+ self._load_geoms()
+ self._load_lights()
+ self._load_cameras()
- if geom.type == USDGeomType.Mesh.value:
- self.usd_geoms.append(USDMesh(self.model.geom_dataid[geom.objid],
- geom,
- self.stage,
- self.model,
- texture_file))
- else:
- self.usd_geoms.append(create_usd_geom_primitive(geom,
- self.stage,
- texture_file))
+ self._update_geoms()
+ self._update_lights()
+ self._update_cameras()
- # initializes an array to store all the lights in the scene
- # populates with "empty" USDLight objects
- self.usd_lights = []
- lights = self.scene.lights
- self.nlight = self.scene.nlight
- for i in range(self.nlight):
- self.usd_lights.append(USDLight(self.stage))
+ self.updates += 1
- # initializes an array to store all the cameras in the scene
- # populates with "empty" USDCamera objects
- self.usd_cameras = []
- ncam = self.model.ncam
- for i in range(ncam):
- self.usd_cameras.append(USDCamera(self.stage))
+ def _load_textures(self):
+ # TODO: remove code once added internally to mujoco
+ data_adr = 0
+ self.texture_files = []
+ for texture_id in range(self.model.ntex):
+ texture_height = self.model.tex_height[texture_id]
+ texture_width = self.model.tex_width[texture_id]
+ pixels = 3*texture_height*texture_width
+ img = im.fromarray(self.model.tex_rgb[data_adr:data_adr+pixels].reshape(texture_height, texture_width, 3))
+ img = ImageOps.flip(img)
- def _update(self):
- self._update_geoms()
- self._update_lights()
- self._update_cameras()
+ texture_file_name = f"texture_{texture_id}.png"
- def _update_geoms(self):
- """
- Updates the geoms to match the current scene
- """
- geoms = self.scene.geoms
- for i in range(self.ngeom):
- if self.usd_geoms[i]: # TODO: remove this once all primitives are added
- self.usd_geoms[i].update_geom(geoms[i])
+ img.save(os.path.join(self.assets_directory, texture_file_name))
- def _update_lights(self):
- """
- Updates the lights to match the current scene
- """
- lights = self.scene.lights
- nlight = self.scene.nlight
- for i in range(nlight):
- self.usd_lights[i].update_light(lights[i])
-
- def _update_cameras(self):
- """
- Updates the camera to match the current scene
- """
- ncam = self.model.ncam
- for i in range(ncam):
- self.usd_cameras[i].update_camera(self.model.cam_pos[i], self.model.cam_quat[i])
+ relative_path = os.path.relpath(self.assets_directory, self.frames_directory)
+ # img_path = os.path.join(relative_path, texture_file_name) # relative path, TODO: switch back to this
- def compress(self):
- """
- Compresses the output directory to a zip file for easy transfer
- """
- if self.verbose:
- output_dir_msg = colored(f"Compressing files at {self.output_dir} and saving at {self.output_dir}", "green")
- print(output_dir_msg)
- shutil.make_archive(base_name=self.output_dir,
- format='zip',
- base_dir=self.root_dir_name)
+ # absolute path for cluster, TODO: remove!
+ abs_path = os.path.join(os.path.abspath(self.assets_directory), texture_file_name)
- def start_viewer(self):
- if self.data:
- viewer.launch(self.model)
+ self.texture_files.append(abs_path)
- def render(self):
- # should render the usd file given a particular renderer that
- # works with USD files
- # TODO: determine if this is valid functionality
- pass
+ data_adr += pixels
- # TODO: remove later, this is only for debugging purposes
- def print_geom_information(self):
- for i in range(self.ngeom):
- print(self.usd_geoms[i])
\ No newline at end of file
+ if self.verbose:
+ print(colored(f"Writing texture {texture_id}", "cyan"))
+
+ if self.verbose:
+ print(colored(f"Completed writing {self.model.ntex} textures to {self.assets_directory}", "green"))
+
+ def _load_geoms(self):
+ # stores a list of all the geoms in the scene
+ self.usd_geoms = []
+
+ # initializing the geoms
+ for i in range(self.scene.ngeom):
+ geom = self.scene.geoms[i]
+
+ if geom.rgba[3] <= 0:
+ self.usd_geoms.append(None)
+ continue
+
+ # handles meshes in scene
+ if geom.type == mjtGeom.mjGEOM_MESH:
+ usd_geom = USDMesh(stage=self.stage,
+ model=self.model,
+ geom=geom,
+ objid=i,
+ dataid=self.model.geom_dataid[geom.objid],
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ # handles primitives
+ else:
+ if geom.type == mjtGeom.mjGEOM_PLANE:
+ usd_geom = USDPlaneMesh(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_SPHERE:
+ usd_geom = USDSphereMesh(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_CAPSULE:
+ usd_geom = USDCapsule(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_ELLIPSOID:
+ usd_geom = USDEllipsoid(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_CYLINDER:
+ usd_geom = USDCylinderMesh(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ elif geom.type == mjtGeom.mjGEOM_BOX:
+ usd_geom = USDCubeMesh(stage=self.stage,
+ geom=geom,
+ objid=i,
+ rgba=geom.rgba,
+ texture_file=self.texture_files[geom.texid] if geom.texid != -1 else None)
+ else:
+ usd_geom = None
+ self.usd_geoms.append(usd_geom)
+
+ def _update_geoms(self):
+ # iterate through all geoms in the scene and makes update
+ for i in range(self.scene.ngeom):
+ geom = self.scene.geoms[i]
+ if self.usd_geoms[i]:
+ self.usd_geoms[i].update(pos=geom.pos,
+ mat=geom.mat,
+ frame=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]
+ self.usd_lights.append(USDLight(stage=self.stage,
+ objid=i))
+
+ def _update_lights(self):
+ for i in range(self.scene.nlight):
+ light = self.scene.lights[i]
+ self.usd_lights[i].update(pos=light.pos,
+ intensity=self.light_intensity,
+ color=light.diffuse,
+ frame=self.updates)
+
+ def _load_cameras(self):
+ self.camera = USDCamera(stage=self.stage,
+ objid=0)
+
+ def _update_cameras(self):
+ camera = mjv_averageCamera(self.scene.camera[0], self.scene.camera[1])
+
+ forward = camera.forward
+ up = camera.up
+ right = np.cross(forward, up)
+
+ R = np.eye(3)
+ R[:, 0] = right
+ R[:, 1] = up
+ R[:, 2] = -forward
+
+ self.camera.update(cam_pos=camera.pos,
+ cam_mat=R,
+ frame=self.updates)
+
+ def add_light(self,
+ pos: List[float],
+ intensity:int,
+ radius: Optional[float] = 1.0,
+ color: Optional[np.array] = np.array([0.3, 0.3, 0.3]),
+ objid: Optional[int]=1):
+ new_light = USDLight(stage=self.stage,
+ objid=objid,
+ radius=radius)
+
+ new_light.update(pos=pos,
+ intensity=intensity,
+ color=color,
+ frame=0)
+
+ def add_camera(self,
+ pos:List[float],
+ rotation_xyz:List[float],
+ objid: Optional[int]=1):
+ # TODO: change this!
+ new_camera = USDCamera(stage=self.stage,
+ objid=objid)
+
+ r = R.from_euler('xyz', rotation_xyz, degrees=True)
+ new_camera.update(cam_pos=pos,
+ cam_mat=r.as_matrix(),
+ frame=0)
+
+ def save_scene(self):
+ self.stage.SetEndTimeCode(self.frame_count)
+ # with open(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd', "w") as f:
+ # f.write(self.usd)
+ self.stage.Export(f'./{self.output_directory_name}/frames/frame_{self.frame_count}_.usd')
+ if self.verbose:
+ print(colored(f"Writing frame_{self.frame_count}", "green"))#
\ No newline at end of file
diff --git a/python/mujoco/usd_utilities.py b/python/mujoco/usd_utilities.py
deleted file mode 100644
index 19ee07f6..00000000
--- a/python/mujoco/usd_utilities.py
+++ /dev/null
@@ -1,16 +0,0 @@
-
-def get_mesh_ranges(nmesh, arr):
- mesh_ranges = [0]
- running_sum = 0
- for i in range(nmesh):
- running_sum += arr[i]
- mesh_ranges.append(running_sum)
- return mesh_ranges
-
-def get_facetexcoord_ranges(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
\ No newline at end of file
diff --git a/python/mujoco/usd_utils.py b/python/mujoco/usd_utils.py
new file mode 100644
index 00000000..4f9d2ccb
--- /dev/null
+++ b/python/mujoco/usd_utils.py
@@ -0,0 +1,12 @@
+import numpy as np
+
+def create_transform_matrix(rotation_matrix, translation_vector):
+ # Ensure rotation_matrix and translation_vector are NumPy arrays
+ rotation_matrix = np.array(rotation_matrix)
+ translation_vector = np.array(translation_vector)
+
+ transform_matrix = np.eye(4)
+ transform_matrix[:3, :3] = rotation_matrix
+ transform_matrix[:3, 3] = translation_vector
+
+ return transform_matrix
\ No newline at end of file