isaac/robozon_sorter/sim/scene.py

"""Loads the real sorting cell (scene/sorter.usd) and applies the runtime configuration
it needs to actually run.

The scene file is the original build - conveyor art, diverters, camera portal, camera
bodies, laser gate and collection bin exactly as authored. Nothing here rebuilds geometry.
What this module does is re-apply the handful of runtime settings that USD does not carry
and that the cell does not work without; each one is documented where it is applied,
because every one of them was a silent failure at some point.
"""
from __future__ import annotations

import json
from pathlib import Path

from pxr import Gf, PhysxSchema, UsdGeom, UsdPhysics, UsdShade

from .. import config as C

SCENE = C.ROOT / "scene" / "sorter.usd"

# --- prim paths in the authored scene -------------------------------------------------
BELTS = ["/World/ConveyorTrack/Belt", "/World/ConveyorTrack_02/Belt",
         "/World/ConveyorTrack_03/Belt", "/World/ConveyorTrack_04/Belt",
         "/World/ConveyorTrack_01/Belt"]
SPAWN_BELT = "/World/SortingRig/SpawnBelt"
BRANCH = "/World/ConveyorTrack_03/Belt_01"          # the branch the pusher feeds
BLADE = "/World/Diverters/DiverterY_Split/Pusher"
PUSHER_JOINT = "/World/Diverters/DiverterY_Split/PusherSlide"
ANIM_GRAPH = "/World/Diverters/DiverterAnimGraph"
CAMERA_BODIES = "/World/CameraBodies"
ITEMS_ROOT = "/World/Items"
RIG = "/RigRS"

# The blade's parent carries this offset; world_y = PARENT_Y + local_y.
BLADE_PARENT_Y = -0.35


def open_scene(usd_path: str | Path | None = None):
    """open sorter.usd into the current context"""
    import omni.usd
    path = str(usd_path or SCENE)
    if not Path(path).exists():
        raise FileNotFoundError(
            f"{path} not found. The conveyor art it references lives in assets/conveyors/ - "
            "run scripts/fetch_assets.py if that folder is empty."
        )
    omni.usd.get_context().open_stage(path)
    return omni.usd.get_context().get_stage()


# --------------------------------------------------------------------- runtime config
def configure_physics(stage):
    scene = stage.GetPrimAtPath("/World/PhysicsScene")
    if not scene.IsValid():
        scene = UsdPhysics.Scene.Define(stage, "/World/PhysicsScene").GetPrim()
        UsdPhysics.Scene(scene).CreateGravityMagnitudeAttr().Set(9.81)
    px = PhysxSchema.PhysxSceneAPI.Apply(scene)
    # 120 Hz is what the cell was tuned and validated at, together with the 2.5 m/s blade.
    # Raising it changes the contact response and the pushed item stops landing in the bin,
    # so treat this number and PUSHER_SPEED as a matched pair.
    px.CreateTimeStepsPerSecondAttr().Set(120)
    px.CreateEnableCCDAttr().Set(True)
    px.CreateSolverTypeAttr().Set("TGS")


def configure_belts(stage, speed=None, grip_path="/World/SortingRig/M_beltPhysics"):
    """explicit surface velocities; the authored ConveyorBeltGraphs carry no speed and
    would only fight these, so they are switched off.

    `grip_path` is where the belt friction material is authored. It defaults to a prim
    under the sorter's rig; plow_cell.usd has no SortingRig and passes its own path so the
    scene does not grow an empty one.
    """
    speed = speed if speed is not None else C.BELT_SPEED

    grip = stage.GetPrimAtPath(grip_path)
    if not grip.IsValid():
        grip = stage.DefinePrim(grip_path, "Material")
        pm = UsdPhysics.MaterialAPI.Apply(grip)
        pm.CreateStaticFrictionAttr().Set(1.1)
        pm.CreateDynamicFrictionAttr().Set(0.95)
        pm.CreateRestitutionAttr().Set(0.02)
    grip_mat = UsdShade.Material(grip)

    def drive(path, vel):
        prim = stage.GetPrimAtPath(path)
        if not prim.IsValid():
            return False
        if not prim.HasAPI(UsdPhysics.RigidBodyAPI):
            UsdPhysics.RigidBodyAPI.Apply(prim)
            UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(True)
        PhysxSchema.PhysxSurfaceVelocityAPI.Apply(prim)
        PhysxSchema.PhysxSurfaceVelocityAPI(prim).CreateSurfaceVelocityAttr().Set(Gf.Vec3f(*vel))
        api = UsdShade.MaterialBindingAPI.Apply(prim)
        api.Bind(grip_mat, bindingStrength=UsdShade.Tokens.strongerThanDescendants,
                 materialPurpose="physics")
        return True

    for path in BELTS + [SPAWN_BELT]:
        drive(path, (-speed, 0, 0))

    # The branch is rotated: its LOCAL X points along world -Y. surfaceVelocity is given
    # in the body's local frame, so carrying goods toward the bin (+Y) needs (-speed,0,0).
    # Setting the "obvious" (0,+speed,0) drags them sideways and they sit there.
    drive(BRANCH, (-speed, 0, 0))

    for track in ["ConveyorTrack", "ConveyorTrack_02", "ConveyorTrack_03",
                  "ConveyorTrack_04", "ConveyorTrack_01"]:
        for graph in (f"/World/{track}/ConveyorBeltGraph", f"/World/{track}/ConveyorBeltGraph_01"):
            g = stage.GetPrimAtPath(graph)
            if g.IsValid():
                g.SetActive(False)


def configure_pusher(stage):
    """the blade is driven kinematically from script.

    Its authored PhysicsPrismaticJoint is unusable at runtime: USD drive-target writes
    reach PhysX about a second late, so the blade never completes its stroke while the
    item is still in reach. The joint is disabled and the blade is moved directly.
    """
    blade = stage.GetPrimAtPath(BLADE)
    if not blade.IsValid():
        raise RuntimeError(f"{BLADE} missing - is this the right scene?")
    UsdPhysics.RigidBodyAPI(blade).CreateKinematicEnabledAttr().Set(True)

    joint = stage.GetPrimAtPath(PUSHER_JOINT)
    if joint.IsValid():
        joint.GetAttribute("physics:jointEnabled").Set(False)

    graph = stage.GetPrimAtPath(ANIM_GRAPH)
    if graph.IsValid():
        graph.SetActive(False)      # otherwise it rewrites the diverter targets every tick

    # the blade must sweep through the conveyor rails rather than grind on them
    filt = UsdPhysics.FilteredPairsAPI.Apply(blade)
    rel = filt.CreateFilteredPairsRel()
    have = {str(t) for t in rel.GetTargets()}
    for path in BELTS + [SPAWN_BELT, BRANCH, "/World/Diverters/DiverterY_Split/Base"]:
        if stage.GetPrimAtPath(path).IsValid() and path not in have:
            rel.AddTarget(path)

    # seat the blade just over the belt so flat items cannot slip underneath
    for op in UsdGeom.Xformable(blade).GetOrderedXformOps():
        if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
            v = op.Get()
            op.Set(Gf.Vec3d(v[0], C.BLADE_HOME_Y - BLADE_PARENT_Y, -0.135))
            break


def hide_aim_markers(stage):
    """the camera bodies carry cosmetic aim-ray cones that sit right over the inspection
    point; left visible they dominate the frame and segmentation locks onto them."""
    n = 0
    for rig in ["RealSense_D435", "Orbbec_Gemini305", "Orbbec_Gemini345"]:
        p = stage.GetPrimAtPath(f"{CAMERA_BODIES}/{rig}/AimRay")
        if p.IsValid():
            UsdGeom.Imageable(p).MakeInvisible()
            n += 1
    return n


def load_test_items(stage, meshes_dir=None):
    """add the bundled per-class test meshes as dynamic rigid bodies"""
    meshes_dir = Path(meshes_dir or C.MESHES)
    manifest = json.loads((meshes_dir / "manifest.json").read_text())
    UsdGeom.Xform.Define(stage, ITEMS_ROOT)
    items = {}
    for i, (name, meta) in enumerate(sorted(manifest.items())):
        usd = meshes_dir / f"{name}.usd"
        if not usd.exists():
            continue
        prim = UsdGeom.Xform.Define(stage, f"{ITEMS_ROOT}/{name}").GetPrim()
        refs = prim.GetReferences()
        refs.ClearReferences()          # idempotent: prepare() may run more than once
        refs.AddReference(str(usd))
        xf = UsdGeom.Xformable(prim)
        xf.ClearXformOpOrder()
        xf.AddTranslateOp().Set(Gf.Vec3d(9.0 + 1.2 * i, 5.0, 0.4))
        UsdPhysics.RigidBodyAPI.Apply(prim)
        # meshes exported from a streaming scene arrive kinematic and hidden - both make
        # them inert: kinematic ignores gravity and belt friction, hidden shows nothing
        UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(False)
        UsdPhysics.MassAPI.Apply(prim).CreateMassAttr().Set(0.6)
        px = PhysxSchema.PhysxRigidBodyAPI.Apply(prim)
        px.CreateEnableCCDAttr().Set(True)
        px.CreateSolverPositionIterationCountAttr().Set(24)
        px.CreateSleepThresholdAttr().Set(0.0)      # a settled item must still be draggable
        UsdGeom.Imageable(prim).MakeVisible()
        items[name] = meta
    return items


def prepare(stage, belt_speed=None, meshes_dir=None):
    """everything the authored scene needs before it will run"""
    configure_physics(stage)
    configure_belts(stage, belt_speed)
    configure_pusher(stage)
    hidden = hide_aim_markers(stage)
    items = load_test_items(stage, meshes_dir)
    calib = json.loads((C.CONFIG / "calib.json").read_text())
    return dict(items=items, calib=calib, aim_markers_hidden=hidden)


def load(usd_path=None, belt_speed=None, meshes_dir=None):
    stage = open_scene(usd_path)
    return stage, prepare(stage, belt_speed, meshes_dir)