USD integration with support for initial primitives

python/mujoco/usd_component.py

@@ -0,0 +1,141 @@
+from enum import Enum
+
+import mujoco
+from pxr import Usd, UsdGeom, Vt
+
+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
+  Sphere = 2
+  Cube = 6
+
+class USDGeom(object):
+  """
+  Parent class for created geoms
+  """
+  def __init__(self,
+               geom=None,
+               stage=None):
+    self.geom = geom
+    self.stage = stage # TODO: remove, not being used
+    self.type = None
+    self.xform = None
+    self.prim = None
+    self.ref = None
+
+  def update_geom(self, new_geom):
+    raise NotImplementedError
+  
+  def update_pos(self, new_pos):
+    pos = tuple([float(x) for x in new_pos])
+    self.xform.AddTranslateOp().Set(pos)
+
+  def update_size(self, new_size):
+    size = tuple([float(x) for x in new_size])
+    self.xform.AddScaleOp().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])]
+    color = self.prim.GetDisplayColorAttr()
+    color.Set(rgba)
+
+  # TODO: create another method for transparency
+  def update_transparency(self, new_transparency):
+    pass
+
+  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):
+    super().__init__(geom, stage)
+    self.type = 0
+    USDPlane.plane_count += 1
+    xform_path = f'/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.Plane.Define(stage, plane_path)
+    self.ref = stage.GetPrimAtPath(plane_path)
+
+  def update_geom(self, new_geom):
+    self.update_pos(new_geom.pos)
+    self.update_size(new_geom.size)
+    self.update_color(new_geom.rgba)
+
+class USDSphere(USDGeom):
+  """
+  Stores information regarding a sphere geom in USD
+  """
+
+  sphere_count = 0
+
+  def __init__(self, 
+               geom=None,
+               stage=None):
+    super().__init__(geom, stage)
+    self.type = 2
+    USDSphere.sphere_count += 1
+    xform_path = f'/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)
+
+  def update_geom(self, new_geom):
+    self.update_pos(new_geom.pos)
+    self.update_size(new_geom.size)
+    self.update_color(new_geom.rgba)
+
+class USDCube(USDGeom):
+  """
+  Stores information regarding a cube geom in USD
+  """
+
+  cube_count = 0
+
+  def __init__(self, 
+               geom=None,
+               stage=None):
+    super().__init__(geom, stage)
+    self.type = 6
+    USDCube.cube_count += 1
+    xform_path = f'/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)
+
+  def update_geom(self, new_geom):
+    self.update_pos(new_geom.pos)
+    self.update_size(new_geom.size)
+    self.update_color(new_geom.rgba)
+
+def create_usd_geom(geom, stage):
+  geom_type = geom.type
+  if geom_type==USDGeomType.Plane.value:
+    return USDPlane(geom, stage)
+  elif geom_type==USDGeomType.Sphere.value:
+    return USDSphere(geom, stage)
+  elif geom_type==USDGeomType.Cube.value:
+    return USDCube(geom, stage)
+  
+class USDLight(object):
+  def __init__(self):
+    pass
+
+  def update_light(self, new_light):
+    pass
+  

python/mujoco/usd_renderer.py

@@ -0,0 +1,104 @@
+
+import mujoco
+import mujoco.viewer as viewer
+from mujoco.usd_component import *
+from pxr import Usd, UsdGeom
+
+class USDRenderer(object):
+  """
+  Renderer class the creates USD representations for mujoco scenes
+  """
+  def __init__(self,
+               model,
+               height=480,
+               width=480):
+    self.model = model
+    self.data = None
+    self.renderer = mujoco.Renderer(model, height, width)
+
+    self.loaded_scene_info = False
+
+    self.stage = Usd.Stage.CreateNew('usd_stage.usda')
+
+  @property
+  def usd(self):
+    return self.stage.GetRootLayer().ExportToString()
+  
+  @property
+  def scene(self):
+    return self.renderer.scene
+  
+  def save_scene(self):
+    self.stage.GetRootLayer().Save()
+
+  def update_scene(self, data):
+    self.renderer.update_scene(data)
+    self.data = data
+
+    if not self.loaded_scene_info:
+      # loads the initial geoms, lights, and camera information 
+      # from the scene
+      self._load()
+      self.loaded_scene_info = True
+    
+    self._update()
+
+  def _load(self):
+    """
+    Loads and initializes the necessary objects to render the scene
+    """
+    # 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):
+      self.usd_geoms.append(create_usd_geom(geoms[i], self.stage))
+
+    # 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())
+
+  def _update(self):
+    self._update_geoms()
+    self._update_lights()
+    self._update_camera()
+
+  def _update_geoms(self):
+    """
+    Updates the geoms to match the current scene
+    """
+    geoms = self.scene.geoms
+    for i in range(self.ngeom):
+      self.usd_geoms[i].update_geom(geoms[i])
+
+  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):
+      print(self.usd_lights[i])
+      
+  def _update_camera(self):
+    pass
+
+  def start_viewer(self):
+    if self.data:
+      viewer.launch(self.model, self.data)
+
+  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
+
+  # TODO: remove later, this is only for debuggin purposes
+  def print_geom_information(self):
+    for i in range(self.ngeom):
+      print(self.usd_geoms[i])