Removing mujoco renderer dependency
This commit is contained in:
@@ -15,9 +15,10 @@
|
|||||||
"""USD exporter."""
|
"""USD exporter."""
|
||||||
|
|
||||||
import os
|
import os
|
||||||
from typing import List, Optional
|
from typing import List, Optional, Union
|
||||||
|
|
||||||
import mujoco
|
import mujoco
|
||||||
|
from mujoco import _structs, _enums, _functions
|
||||||
import mujoco.usd.camera as camera_module
|
import mujoco.usd.camera as camera_module
|
||||||
import mujoco.usd.lights as light_module
|
import mujoco.usd.lights as light_module
|
||||||
import mujoco.usd.objects as object_module
|
import mujoco.usd.objects as object_module
|
||||||
@@ -113,7 +114,7 @@ class USDExporter:
|
|||||||
self.usd_cameras = []
|
self.usd_cameras = []
|
||||||
|
|
||||||
# initializing rendering requirements
|
# initializing rendering requirements
|
||||||
self.renderer = mujoco.Renderer(model, height, width, max_geom)
|
self._scene = _structs.MjvScene(model=model, maxgeom=max_geom)
|
||||||
self._initialize_usd_stage()
|
self._initialize_usd_stage()
|
||||||
self._scene_option = mujoco.MjvOption() # using default scene option
|
self._scene_option = mujoco.MjvOption() # using default scene option
|
||||||
|
|
||||||
@@ -131,7 +132,7 @@ class USDExporter:
|
|||||||
@property
|
@property
|
||||||
def scene(self):
|
def scene(self):
|
||||||
"""Returns the scene."""
|
"""Returns the scene."""
|
||||||
return self.renderer.scene
|
return self._scene
|
||||||
|
|
||||||
def _initialize_usd_stage(self):
|
def _initialize_usd_stage(self):
|
||||||
"""Initializes a USD stage to represent the mujoco scene."""
|
"""Initializes a USD stage to represent the mujoco scene."""
|
||||||
@@ -164,10 +165,48 @@ class USDExporter:
|
|||||||
|
|
||||||
if self.verbose:
|
if self.verbose:
|
||||||
print("Writing output frames and assets to"
|
print("Writing output frames and assets to"
|
||||||
f" {self.output_directory_path}",
|
f" {self.output_directory_path}"
|
||||||
"green"
|
|
||||||
)
|
)
|
||||||
|
|
||||||
|
def _update_scene(
|
||||||
|
self,
|
||||||
|
data: _structs.MjData,
|
||||||
|
camera: Union[int, str] = -1,
|
||||||
|
scene_option: Optional[_structs.MjvOption] = None
|
||||||
|
):
|
||||||
|
camera_id = camera
|
||||||
|
if isinstance(camera_id, str):
|
||||||
|
camera_id = _functions.mj_name2id(
|
||||||
|
self.model, _enums.mjtObj.mjOBJ_CAMERA.value, camera_id
|
||||||
|
)
|
||||||
|
if camera_id == -1:
|
||||||
|
raise ValueError(f"The camera '{camera}' does not exist.")
|
||||||
|
if camera_id < -1 or camera_id >= self.model.ncam:
|
||||||
|
raise ValueError(f"The camera id {camera_id} is out of"
|
||||||
|
f" range [-1, {self.model.ncam}).")
|
||||||
|
|
||||||
|
# Render camera.
|
||||||
|
camera = _structs.MjvCamera()
|
||||||
|
camera.fixedcamid = camera_id
|
||||||
|
|
||||||
|
# Defaults to mjCAMERA_FREE, otherwise mjCAMERA_FIXED refers to a
|
||||||
|
# camera explicitly defined in the model.
|
||||||
|
if camera_id == -1:
|
||||||
|
camera.type = _enums.mjtCamera.mjCAMERA_FREE
|
||||||
|
_functions.mjv_defaultFreeCamera(self.model, camera)
|
||||||
|
else:
|
||||||
|
camera.type = _enums.mjtCamera.mjCAMERA_FIXED
|
||||||
|
|
||||||
|
scene_option = scene_option or self._scene_option
|
||||||
|
_functions.mjv_updateScene(
|
||||||
|
self.model,
|
||||||
|
data,
|
||||||
|
scene_option,
|
||||||
|
None,
|
||||||
|
camera, _enums.mjtCatBit.mjCAT_ALL.value,
|
||||||
|
self._scene,
|
||||||
|
)
|
||||||
|
|
||||||
def update_scene(
|
def update_scene(
|
||||||
self,
|
self,
|
||||||
data: mujoco.MjData,
|
data: mujoco.MjData,
|
||||||
@@ -185,7 +224,7 @@ class USDExporter:
|
|||||||
scene_option = scene_option or self._scene_option
|
scene_option = scene_option or self._scene_option
|
||||||
|
|
||||||
# update the mujoco renderer
|
# update the mujoco renderer
|
||||||
self.renderer.update_scene(data, scene_option=scene_option)
|
self._update_scene(data, scene_option=scene_option)
|
||||||
|
|
||||||
if self.updates == 0:
|
if self.updates == 0:
|
||||||
self._initialize_usd_stage()
|
self._initialize_usd_stage()
|
||||||
@@ -232,7 +271,6 @@ class USDExporter:
|
|||||||
if self.verbose:
|
if self.verbose:
|
||||||
print(f"Completed writing {self.model.ntex} textures to"
|
print(f"Completed writing {self.model.ntex} textures to"
|
||||||
f" {self.assets_directory}",
|
f" {self.assets_directory}",
|
||||||
"green",
|
|
||||||
)
|
)
|
||||||
|
|
||||||
def _load_geom(self, geom: mujoco.MjvGeom):
|
def _load_geom(self, geom: mujoco.MjvGeom):
|
||||||
@@ -374,8 +412,8 @@ class USDExporter:
|
|||||||
camera = self.usd_cameras[i]
|
camera = self.usd_cameras[i]
|
||||||
camera_name = self.camera_names[i]
|
camera_name = self.camera_names[i]
|
||||||
|
|
||||||
self.renderer.update_scene(
|
self._update_scene(
|
||||||
data, scene_option=scene_option, camera=camera_name
|
data, camera=camera_name, scene_option=scene_option,
|
||||||
)
|
)
|
||||||
|
|
||||||
avg_camera = mujoco.mjv_averageCamera(
|
avg_camera = mujoco.mjv_averageCamera(
|
||||||
@@ -441,11 +479,11 @@ class USDExporter:
|
|||||||
new_camera = camera_module.USDCamera(
|
new_camera = camera_module.USDCamera(
|
||||||
stage=self.stage, obj_name=obj_name)
|
stage=self.stage, obj_name=obj_name)
|
||||||
|
|
||||||
R = np.zeros(9)
|
rotation = np.zeros(9)
|
||||||
quat = np.zeros(4)
|
quat = np.zeros(4)
|
||||||
mujoco.mju_euler2Quat(quat, rotation_xyz, "xyz")
|
mujoco.mju_euler2Quat(quat, rotation_xyz, "xyz")
|
||||||
mujoco.mju_quat2Mat(R, quat)
|
mujoco.mju_quat2Mat(rotation, quat)
|
||||||
new_camera.update(cam_pos=np.array(pos), cam_mat=R, frame=0)
|
new_camera.update(cam_pos=np.array(pos), cam_mat=rotation, frame=0)
|
||||||
|
|
||||||
def save_scene(self, filetype: str = "usd"):
|
def save_scene(self, filetype: str = "usd"):
|
||||||
"""Saves the scene to a USD file."""
|
"""Saves the scene to a USD file."""
|
||||||
@@ -462,7 +500,7 @@ class USDExporter:
|
|||||||
)
|
)
|
||||||
if self.verbose:
|
if self.verbose:
|
||||||
print(
|
print(
|
||||||
f"Completed writing frame_{self.frame_count}.{filetype}", "green"
|
f"Completed writing frame_{self.frame_count}.{filetype}"
|
||||||
)
|
)
|
||||||
|
|
||||||
def _get_geom_name(self, geom):
|
def _get_geom_name(self, geom):
|
||||||
|
|||||||
Reference in New Issue
Block a user