diff --git a/python/dist/mujoco-2.3.8.tar.gz b/python/dist/mujoco-2.3.8.tar.gz new file mode 100644 index 00000000..9f101d96 Binary files /dev/null and b/python/dist/mujoco-2.3.8.tar.gz differ diff --git a/python/mujoco/usd_component.py b/python/mujoco/usd_component.py new file mode 100644 index 00000000..7dee1993 --- /dev/null +++ b/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 + diff --git a/python/mujoco/usd_renderer.py b/python/mujoco/usd_renderer.py new file mode 100644 index 00000000..d17b9ef5 --- /dev/null +++ b/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]) \ No newline at end of file