chenlin/Mujoco_WASM
1
0
Fork 0
Code Issues Pull Requests Actions 2 Packages Projects Releases Wiki Activity

Merge pull request #1714 from abhihjoshi:tendons

Browse Source 
PiperOrigin-RevId: 651048034
Change-Id: If5b7cb55ff0007c3c6f932e62f068d1168778a33
This commit is contained in:
Copybara-Service
2024-07-10 09:26:54 -07:00
parent 7f9f8b5429 e734527e7b
commit 60afda949c
9 changed files with 1092 additions and 969 deletions
Download Patch File Download Diff File Expand all files Collapse all files
python/mujoco/usd/exporter.py
+167 -109
View File
@@ -12,27 +12,30 @@
# See the License for the specific language governing permissions and
# limitations under the License.
# ==============================================================================
"""USD exporter."""
import os
from typing import List, Optional
import mujoco
import mujoco.usd.component as component_module
import mujoco.usd.camera as camera_module
import mujoco.usd.lights as light_module
import mujoco.usd.objects as object_module
import mujoco.usd.shapes as shapes_module
import numpy as np
from PIL import Image as im
from PIL import ImageOps
import scipy
import termcolor
import tqdm
from typing import List, Optional, Tuple, Union
from PIL import Image as im
from PIL import ImageOps
# TODO: b/288149332 - Remove once USD Python Binding works well with pytype.
# pytype: disable=module-attr
from pxr import Sdf
from pxr import Usd
from pxr import UsdGeom
class USDExporter:
"""MuJoCo to USD exporter for porting scenes to external renderers."""
def __init__(
self,
@@ -47,13 +50,13 @@ class USDExporter:
specialized_materials_file: Optional[str] = None,
verbose: bool = True,
):
"""Initializes a new USD Exporter
"""Initializes a new USD Exporter.
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
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.
@@ -61,6 +64,10 @@ class USDExporter:
and assets generated by the USD renderer.
output_directory_root: path to root directory storing generated frames
and assets by the USD renderer.
light_intensity: intensity of the light in the scene.
camera_names: list of camera names to be used in the scene.
specialized_materials_file: path to a file containing a list of
materials to be used in the scene.
verbose: decides whether to print updates.
"""
@@ -99,7 +106,12 @@ class USDExporter:
self.frame_count = 0 # maintains how many times we have saved the scene
self.updates = 0
self.geom_name2usd = {}
self.geom_names = set()
self.geom_refs = {}
# initializing list of lights and cameras
self.usd_lights = []
self.usd_cameras = []
# initializing rendering requirements
self.renderer = mujoco.Renderer(model, height, width, max_geom)
@@ -114,18 +126,21 @@ class USDExporter:
@property
def usd(self):
"""Returns the USD file as a string."""
return self.stage.GetRootLayer().ExportToString()
@property
def scene(self):
"""Returns the scene."""
return self.renderer.scene
def _initialize_usd_stage(self):
"""Initializes a USD stage to represent the mujoco scene."""
self.stage = Usd.Stage.CreateInMemory()
UsdGeom.SetStageUpAxis(self.stage, UsdGeom.Tokens.z)
self.stage.SetStartTimeCode(0)
# add as user input
self.stage.SetTimeCodesPerSecond(24.0)
self.stage.SetTimeCodesPerSecond(60.0)
default_prim = UsdGeom.Xform.Define(
self.stage, Sdf.Path("/World")
@@ -133,6 +148,7 @@ class USDExporter:
self.stage.SetDefaultPrim(default_prim)
def _initialize_output_directories(self):
"""Initializes output directories to store frames and assets."""
self.output_directory_path = os.path.join(
self.output_directory_root, self.output_directory_name
)
@@ -161,7 +177,7 @@ class USDExporter:
data: mujoco.MjData,
scene_option: Optional[mujoco.MjvOption] = None,
):
"""Updates the scene with latest sim data
"""Updates the scene with latest sim data.
Args:
data: structure storing current simulation state
@@ -175,7 +191,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()
@@ -189,7 +204,7 @@ class USDExporter:
self.updates += 1
def _load_textures(self):
# TODO: remove code once added internally to mujoco
"""Load textures."""
data_adr = 0
self.texture_files = []
for texture_id in tqdm.tqdm(range(self.model.ntex)):
@@ -212,7 +227,7 @@ class USDExporter:
)
img_path = os.path.join(
relative_path, texture_file_name
) # relative path, TODO: switch back to this
)
self.texture_files.append(img_path)
@@ -228,19 +243,20 @@ class USDExporter:
)
def _load_geom(self, geom: mujoco.MjvGeom):
"""Loads a geom into the USD scene."""
geom_name = self._get_geom_name(geom)
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
assert geom_name not in self.geom_name2usd
assert geom_name not in self.geom_names
texture_file = (
self.texture_files[self.model.mat_texid[mujoco.mjNTEXMAT*geom.matid]]
self.texture_files[self.model.mat_texid[geom.matid][0]]
if geom.matid != -1
else None
)
# handles meshes in scene
# 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,
@@ -249,95 +265,77 @@ class USDExporter:
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_PLANE:
usd_geom = component_module.USDPlaneMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_SPHERE:
usd_geom = component_module.USDSphereMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_CAPSULE:
usd_geom = component_module.USDCapsule(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_ELLIPSOID:
usd_geom = component_module.USDEllipsoid(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_CYLINDER:
usd_geom = component_module.USDCylinderMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
elif geom.type == mujoco.mjtGeom.mjGEOM_BOX:
usd_geom = component_module.USDCubeMesh(
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
else:
usd_geom = None
# handling tendons in our scene
if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
mesh_config = shapes_module.mesh_config_generator(
name=geom_name,
geom_type=geom.type,
size=np.array([1.0, 1.0, 1.0]),
decouple=True
)
usd_geom = object_module.USDTendon(
mesh_config=mesh_config,
stage=self.stage,
geom=geom,
obj_name=geom_name,
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 = object_module.USDPrimitiveMesh(
mesh_config=mesh_config,
stage=self.stage,
geom=geom,
obj_name=geom_name,
rgba=geom.rgba,
texture_file=texture_file,
)
self.geom_name2usd[geom_name] = usd_geom
self.geom_names.add(geom_name)
self.geom_refs[geom_name] = usd_geom
def _update_geoms(self):
geom_names = set(self.geom_name2usd.keys())
# iterate through all geoms in the scene and makes update
"""Iterate through all geoms in the scene and makes update."""
for i in range(self.scene.ngeom):
geom = self.scene.geoms[i]
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
geom_name = self._get_geom_name(geom)
if geom_name not in self.geom_name2usd:
if geom_name not in self.geom_names:
# load a new object into USD
self._load_geom(geom)
if self.geom_name2usd[geom_name]:
self.geom_name2usd[geom_name].update_visibility(False, 0)
if self.geom_name2usd[geom_name]:
self.geom_name2usd[geom_name].update(
if geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
tendon_scale = geom.size
self.geom_refs[geom_name].update(
pos=geom.pos,
mat=geom.mat,
scale=tendon_scale,
visible=geom.rgba[3] > 0,
frame=self.updates,
)
else:
self.geom_refs[geom_name].update(
pos=geom.pos,
mat=geom.mat,
visible=geom.rgba[3] > 0,
frame=self.updates,
)
if geom_name in geom_names:
geom_names.remove(geom_name)
for geom_name in geom_names:
if self.geom_name2usd[geom_name]:
self.geom_name2usd[geom_name].update_visibility(False, 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]
if not np.allclose(light.pos, [0, 0, 0]):
self.usd_lights.append
(component_module.USDSphereLight(stage=self.stage, obj_name=str(i)))
self.usd_lights.append(
light_module.USDSphereLight(stage=self.stage, obj_name=str(i))
)
else:
self.usd_lights.append(None)
@@ -359,19 +357,23 @@ class USDExporter:
)
def _load_cameras(self):
self.usd_cameras = []
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,
data: mujoco.MjData,
scene_option: Optional[mujoco.MjvOption] = None,
):
for i in range(len(self.usd_cameras)):
"""Updates cameras.
Args:
data: An MjData instance.
scene_option: An optional MjvOption instance.
"""
for i in range(len(self.usd_cameras)):
camera = self.usd_cameras[i]
camera_name = self.camera_names[i]
@@ -386,12 +388,14 @@ class USDExporter:
up = avg_camera.up
right = np.cross(forward, up)
R = np.eye(3)
R[:, 0] = right
R[:, 1] = up
R[:, 2] = -forward
rotation = np.eye(3)
rotation[:, 0] = right
rotation[:, 1] = up
rotation[:, 2] = -forward
camera.update(cam_pos=avg_camera.pos, cam_mat=R, frame=self.updates)
camera.update(
cam_pos=avg_camera.pos, cam_mat=rotation, frame=self.updates
)
def add_light(
self,
@@ -399,37 +403,91 @@ class USDExporter:
intensity: int,
radius: Optional[float] = 1.0,
color: Optional[np.ndarray] = np.array([0.3, 0.3, 0.3]),
objid: Optional[int] = 1,
obj_name: Optional[str] = "light_1",
light_type: Optional[str] = "sphere",
):
"""Adds a user defined, fixed light.
Args:
pos: position of the light in 3D space.
intensity: intensity of the light.
radius: radius of the light to be used by renderer.
color: color of the light.
obj_name: name associated with the light.
light_type: type of light (sphere or dome).
"""
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=obj_name, radius=radius
)
new_light.update(pos=np.array(pos), intensity=intensity, color=color, frame=0)
new_light.update(
pos=np.array(pos), intensity=intensity, color=color, frame=0
)
elif light_type == "dome":
new_light = component_module.USDDomeLight(
stage=self.stage, obj_name=str(objid))
new_light.update(intensity=intensity, color=color, frame=0)
new_light = light_module.USDDomeLight(stage=self.stage, obj_name=obj_name)
new_light.update(intensity=intensity, color=color)
def add_camera(
self,
pos: List[float],
rotation_xyz: List[float],
objid: Optional[int] = 1,
obj_name: Optional[str] = "camera_1",
):
new_camera = component_module.USDCamera(
stage=self.stage, obj_name=str(objid))
"""Adds a user defined, fixed camera.
Args:
pos: position of the camera in 3D space.
rotation_xyz: euler rotation of the camera.
obj_name: name associated with the camera.
"""
new_camera = camera_module.USDCamera(
stage=self.stage, obj_name=obj_name)
r = scipy.spatial.transform.Rotation.from_euler(
"xyz", rotation_xyz, degrees=True)
new_camera.update(cam_pos=np.array(pos), cam_mat=r.as_matrix(), frame=0)
def save_scene(self, filetype: str = "usd"):
"""Saves the scene to a USD file."""
assert filetype in ["usd", "usda", "usdc"]
self.stage.SetEndTimeCode(self.frame_count)
# post-processing for visibility of geoms in scene
for _, geom_ref in self.geom_refs.items():
geom_ref.update_visibility(False, geom_ref.last_visible_frame+1)
self.stage.Export(
f"{self.output_directory_root}/{self.output_directory_name}/frames/frame_{self.frame_count}_.{filetype}"
f"{self.output_directory_root}/{self.output_directory_name}/" +
f"frames/frame_{self.frame_count}.{filetype}"
)
if self.verbose:
print(termcolor.colored(f"Writing frame_{self.frame_count}", "green"))
print(
termcolor.colored(
f"Completed writing frame_{self.frame_count}.{filetype}", "green"
)
)
def _get_geom_name(self, geom):
"""Adding id as part of name for USD file."""
geom_name = mujoco.mj_id2name(self.model, geom.objtype, geom.objid)
if not geom_name:
geom_name = "None"
geom_name += f"_id{geom.objid}"
# adding additional naming information to differentiate
# between geoms and tendons
if geom.objtype == mujoco.mjtObj.mjOBJ_GEOM:
geom_name += "_geom"
elif geom.objtype == mujoco.mjtObj.mjOBJ_TENDON:
geom_name += f"_tendon_segid{geom.segid}"
return geom_name
# for debugging purposes, prints all geoms in scene
# including those part of tendons
def _print_scene_geom_info(self):
for i in range(self.scene.ngeom):
geom = self.scene.geoms[i]
geom_name = self._get_geom_name(geom)
print(i, geom_name)