Spaces:
Paused
Paused
| """Build an animated USD stage from a MuJoCo scene and a run's trajectory, for rendering in Isaac Sim. | |
| The stage is a kinematic replay: every MuJoCo body becomes an Xform under /World/Bodies whose world pose is | |
| time-sampled at the trajectory's frame rate, with its visible geoms (MuJoCo's viewer groups 0-2) as children. | |
| A camera, time-sampled to follow the robot like the MuJoCo video, sits at /World/Camera. No physics runs in | |
| Isaac Sim; it only renders what MuJoCo simulated, so both videos show exactly the same motion. | |
| Pure `pxr` (usd-core) + mujoco, so it can be built and tested anywhere; Isaac Sim just opens the file. | |
| """ | |
| from __future__ import annotations | |
| import math | |
| from pathlib import Path | |
| import mujoco | |
| import numpy as np | |
| VISIBLE_GROUPS = (0, 1, 2) | |
| PLANE_SIZE = 40.0 # metres, for MuJoCo's infinite planes | |
| def _name(s: str) -> str: | |
| """A valid USD prim name.""" | |
| out = "".join(c if c.isalnum() or c == "_" else "_" for c in s) | |
| return out if out and not out[0].isdigit() else f"_{out}" | |
| def _quat(q): | |
| from pxr import Gf | |
| return Gf.Quatf(float(q[0]), Gf.Vec3f(float(q[1]), float(q[2]), float(q[3]))) | |
| def _material(stage, rgba, cache: dict): | |
| from pxr import Gf, Sdf, UsdShade | |
| key = tuple(round(float(c), 3) for c in rgba) | |
| if key in cache: | |
| return cache[key] | |
| path = f"/World/Looks/m{len(cache)}" | |
| mat = UsdShade.Material.Define(stage, path) | |
| shader = UsdShade.Shader.Define(stage, f"{path}/Shader") | |
| shader.CreateIdAttr("UsdPreviewSurface") | |
| shader.CreateInput("diffuseColor", Sdf.ValueTypeNames.Color3f).Set(Gf.Vec3f(*key[:3])) | |
| shader.CreateInput("opacity", Sdf.ValueTypeNames.Float).Set(key[3]) | |
| shader.CreateInput("roughness", Sdf.ValueTypeNames.Float).Set(0.6) | |
| mat.CreateSurfaceOutput().ConnectToSource(shader.ConnectableAPI(), "surface") | |
| cache[key] = mat | |
| return mat | |
| def _geom_prim(stage, model: mujoco.MjModel, g: int, parent: str): | |
| """One USD prim for MuJoCo geom g, placed in its body's frame. Returns None for types we do not draw.""" | |
| from pxr import Gf, UsdGeom, Vt | |
| T = mujoco.mjtGeom | |
| gtype, size = model.geom_type[g], model.geom_size[g] | |
| path = f"{parent}/g{g}" | |
| if gtype == T.mjGEOM_PLANE: | |
| prim = UsdGeom.Cube.Define(stage, path) | |
| prim.CreateSizeAttr(1.0) | |
| sx = size[0] * 2 if size[0] > 0 else PLANE_SIZE | |
| sy = size[1] * 2 if size[1] > 0 else PLANE_SIZE | |
| scale = Gf.Vec3f(sx, sy, 0.002) | |
| offset = Gf.Vec3d(0, 0, -0.001) | |
| elif gtype == T.mjGEOM_BOX: | |
| prim = UsdGeom.Cube.Define(stage, path) | |
| prim.CreateSizeAttr(2.0) | |
| scale, offset = Gf.Vec3f(*map(float, size)), None | |
| elif gtype in (T.mjGEOM_SPHERE, T.mjGEOM_ELLIPSOID): | |
| prim = UsdGeom.Sphere.Define(stage, path) | |
| prim.CreateRadiusAttr(1.0) | |
| r = size if gtype == T.mjGEOM_ELLIPSOID else (size[0],) * 3 | |
| scale, offset = Gf.Vec3f(*map(float, r)), None | |
| elif gtype == T.mjGEOM_CAPSULE: | |
| prim = UsdGeom.Capsule.Define(stage, path) | |
| prim.CreateRadiusAttr(float(size[0])) | |
| prim.CreateHeightAttr(float(2 * size[1])) | |
| prim.CreateAxisAttr("Z") | |
| scale, offset = None, None | |
| elif gtype == T.mjGEOM_CYLINDER: | |
| prim = UsdGeom.Cylinder.Define(stage, path) | |
| prim.CreateRadiusAttr(float(size[0])) | |
| prim.CreateHeightAttr(float(2 * size[1])) | |
| prim.CreateAxisAttr("Z") | |
| scale, offset = None, None | |
| elif gtype == T.mjGEOM_MESH: | |
| m = int(model.geom_dataid[g]) | |
| va, vn = model.mesh_vertadr[m], model.mesh_vertnum[m] | |
| fa, fn = model.mesh_faceadr[m], model.mesh_facenum[m] | |
| prim = UsdGeom.Mesh.Define(stage, path) | |
| prim.CreatePointsAttr(Vt.Vec3fArray.FromNumpy(model.mesh_vert[va : va + vn].astype(np.float32))) | |
| prim.CreateFaceVertexCountsAttr(Vt.IntArray([3] * int(fn))) | |
| prim.CreateFaceVertexIndicesAttr( | |
| Vt.IntArray.FromNumpy(model.mesh_face[fa : fa + fn].reshape(-1).astype(np.int32)) | |
| ) | |
| prim.CreateSubdivisionSchemeAttr("none") | |
| scale, offset = None, None | |
| else: # heightfields, SDFs: not drawn (they are in the MuJoCo video) | |
| return None | |
| xf = UsdGeom.Xformable(prim) | |
| pos = Gf.Vec3d(*map(float, model.geom_pos[g])) | |
| if offset is not None: | |
| pos += offset | |
| xf.AddTranslateOp().Set(pos) | |
| xf.AddOrientOp().Set(_quat(model.geom_quat[g])) | |
| if scale is not None: | |
| xf.AddScaleOp().Set(scale) | |
| return prim | |
| def _camera_pose(target: np.ndarray, azimuth: float, elevation: float, distance: float): | |
| """Eye position for MuJoCo's free/tracking camera convention (degrees).""" | |
| az, el = math.radians(azimuth), math.radians(elevation) | |
| forward = np.array([math.cos(el) * math.cos(az), math.cos(el) * math.sin(az), math.sin(el)]) | |
| return target - distance * forward | |
| def build_stage( | |
| scene: Path, | |
| trajectory: Path, | |
| out: Path, | |
| root_body: str | None, | |
| camera_distance: float = 4.5, | |
| max_seconds: float | None = None, | |
| ) -> dict: | |
| """Write out (a .usda/.usdc) and return a short description of what is in it.""" | |
| from pxr import Gf, Usd, UsdGeom, UsdLux, UsdShade | |
| model = mujoco.MjModel.from_xml_path(str(scene)) | |
| traj = np.load(trajectory) | |
| fps = int(traj["fps"]) | |
| pos, quat = traj["body_pos"], traj["body_quat"] | |
| names = [str(n) for n in traj["body_names"]] | |
| if len(names) != model.nbody: | |
| raise ValueError("the trajectory does not belong to this scene (body count differs)") | |
| frames = len(pos) if max_seconds is None else min(len(pos), int(max_seconds * fps) + 1) | |
| stage = Usd.Stage.CreateNew(str(out)) # .usdc: binary crate file (much smaller than .usda) | |
| UsdGeom.SetStageUpAxis(stage, UsdGeom.Tokens.z) | |
| UsdGeom.SetStageMetersPerUnit(stage, 1.0) | |
| stage.SetTimeCodesPerSecond(fps) | |
| stage.SetFramesPerSecond(fps) | |
| stage.SetStartTimeCode(0) | |
| stage.SetEndTimeCode(frames - 1) | |
| world = UsdGeom.Xform.Define(stage, "/World") | |
| stage.SetDefaultPrim(world.GetPrim()) | |
| UsdGeom.Scope.Define(stage, "/World/Looks") | |
| UsdGeom.Scope.Define(stage, "/World/Bodies") | |
| materials: dict = {} | |
| drawn = 0 | |
| for b in range(model.nbody): | |
| geoms = [g for g in range(model.ngeom) if model.geom_bodyid[g] == b and model.geom_group[g] in VISIBLE_GROUPS] | |
| if not geoms: | |
| continue | |
| path = f"/World/Bodies/{_name(names[b])}_{b}" | |
| body = UsdGeom.Xform.Define(stage, path) | |
| t_op, o_op = body.AddTranslateOp(), body.AddOrientOp() | |
| if b == 0: # the world body never moves | |
| t_op.Set(Gf.Vec3d(0, 0, 0)) | |
| o_op.Set(Gf.Quatf(1, 0, 0, 0)) | |
| else: | |
| for i in range(frames): | |
| t_op.Set(Gf.Vec3d(*map(float, pos[i, b])), Usd.TimeCode(i)) | |
| o_op.Set(_quat(quat[i, b]), Usd.TimeCode(i)) | |
| for g in geoms: | |
| prim = _geom_prim(stage, model, g, path) | |
| if prim is None: | |
| continue | |
| rgba = model.mat_rgba[model.geom_matid[g]] if model.geom_matid[g] >= 0 else model.geom_rgba[g] | |
| UsdShade.MaterialBindingAPI.Apply(prim.GetPrim()).Bind(_material(stage, rgba, materials)) | |
| drawn += 1 | |
| sun = UsdLux.DistantLight.Define(stage, "/World/Sun") | |
| sun.CreateIntensityAttr(3000) | |
| sun.CreateAngleAttr(1.0) | |
| UsdGeom.Xformable(sun).AddRotateXYZOp().Set(Gf.Vec3f(35, 0, 20)) | |
| dome = UsdLux.DomeLight.Define(stage, "/World/Sky") | |
| dome.CreateIntensityAttr(600) | |
| cam = UsdGeom.Camera.Define(stage, "/World/Camera") | |
| cam.CreateFocalLengthAttr(18.0) | |
| cam.CreateClippingRangeAttr(Gf.Vec2f(0.05, 200.0)) | |
| xf = cam.AddTransformOp() | |
| root = mujoco.mj_name2id(model, mujoco.mjtObj.mjOBJ_BODY, root_body) if root_body else -1 | |
| for i in range(frames): | |
| target = pos[i, root] if root > 0 else np.asarray(model.stat.center) | |
| distance = camera_distance if root > 0 else 1.5 * float(model.stat.extent) | |
| eye = _camera_pose(target, 20.0, -22.0, distance) | |
| view = Gf.Matrix4d().SetLookAt(Gf.Vec3d(*map(float, eye)), Gf.Vec3d(*map(float, target)), Gf.Vec3d(0, 0, 1)) | |
| xf.Set(view.GetInverse(), Usd.TimeCode(i)) | |
| if root <= 0: | |
| break # static camera | |
| stage.GetRootLayer().Save() | |
| return {"frames": frames, "fps": fps, "bodies": model.nbody, "geoms_drawn": drawn, "materials": len(materials)} | |
| __all__ = ["build_stage"] | |