"""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"]