isaac/robozon_sorter/sim/mechanics.py

"""Runtime side of the cell: releasing goods, the laser gate and the pusher stroke.

Two hard-won rules are encoded here and should not be "simplified" away:

1. During simulation, read poses from RigidPrim.get_world_poses(). BBoxCache / XformCache
   return the AUTHORED transform, so a moving item looks frozen and every gate misfires.
2. The blade retracts only when (a) the pushed item has cleared the belt and (b) nothing
   else is inside the blade's footprint. Retracting blindly sweeps the blade back through
   the next item and knocks it over.
"""
from __future__ import annotations

import numpy as np
from pxr import Gf, UsdGeom

from .. import config as C
from .scene import BLADE as BLADE_PATH, ITEMS_ROOT, BLADE_PARENT_Y


class Cell:
    def __init__(self, stage, items):
        self.stage = stage
        self.items = list(items)
        self._blade_op = self._blade_translate_op()
        self._blade_base = self._blade_op.Get()
        from isaacsim.core.experimental.prims import RigidPrim
        self._rp = {n: RigidPrim(paths=[f"{ITEMS_ROOT}/{n}"]) for n in self.items}
        from omni.physx import get_physx_scene_query_interface
        self._query = get_physx_scene_query_interface()
        self.blade_to(C.BLADE_HOME_Y)

    # -- blade --------------------------------------------------------------
    def _blade_translate_op(self):
        prim = self.stage.GetPrimAtPath(BLADE_PATH)
        for op in UsdGeom.Xformable(prim).GetOrderedXformOps():
            if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
                return op
        raise RuntimeError(f"{BLADE_PATH} has no translate op to drive")

    def blade_to(self, y):
        """y is a WORLD coordinate; the op lives in the diverter's frame"""
        b = self._blade_base
        self._blade_op.Set(Gf.Vec3d(b[0], y - BLADE_PARENT_Y, b[2]))

    async def stroke(self, app, out=True, speed=None):
        """sweep the blade at a commanded m/s; fine steps keep the contact impulse sane.
        Above ~2.5 m/s the kinematic blade throws goods off the line."""
        speed = min(speed or C.PUSHER_SPEED, C.PUSHER_MAX_SAFE)
        a, b = (C.BLADE_HOME_Y, C.BLADE_OUT_Y) if out else (C.BLADE_OUT_Y, C.BLADE_HOME_Y)
        dt = 1.0 / 60.0
        steps = max(4, int(round(abs(b - a) / max(speed * dt, 1e-6))))
        for i in range(steps + 1):
            self.blade_to(a + (b - a) * i / steps)
            await app.update_app_async(steps=1)
        return steps * dt

    # -- item state ---------------------------------------------------------
    def pose(self, name):
        return self._rp[name].get_world_poses()[0].numpy()[0]

    def place(self, name, pos):
        prim = self.stage.GetPrimAtPath(f"{ITEMS_ROOT}/{name}")
        for op in UsdGeom.Xformable(prim).GetOrderedXformOps():
            if op.GetOpType() == UsdGeom.XformOp.TypeTranslate:
                op.Set(Gf.Vec3d(*pos))
                return
            if op.GetOpType() == UsdGeom.XformOp.TypeTransform:
                M = Gf.Matrix4d(op.Get())
                M.SetTranslateOnly(Gf.Vec3d(*pos))
                op.Set(M)
                return

    def _underside_gap(self, name):
        """distance from the prim origin down to its lowest point, so it can be seated"""
        cache = UsdGeom.BBoxCache(0, ["default", "render"], useExtentsHint=True)
        prim = self.stage.GetPrimAtPath(f"{ITEMS_ROOT}/{name}")
        r = cache.ComputeWorldBound(prim).ComputeAlignedRange()
        if r.IsEmpty():
            return 0.0
        origin = UsdGeom.Xformable(prim).ComputeLocalToWorldTransform(0).ExtractTranslation()
        return origin[2] - r.GetMin()[2]

    def park(self, name, index=0):
        self.place(name, (9.0 + 1.2 * index, 5.0, 0.4))

    def park_all(self):
        """park everything AND freeze it, so the queue does not fall out of the world.

        Parked items are ordinary dynamic bodies sitting off to the side at y ~ +5, where
        there is no floor under them - so the whole undispatched queue free-falls for the
        entire run. Measured: parked stock at z = -665 after a minute and the stage bound
        reaching z = -20438, which also wrecks every "frame the whole scene" camera because
        the scene is suddenly 20 km tall. Freezing them costs nothing and they are woken in
        `release`, which clears the flag before placing the item on the belt.
        """
        from pxr import UsdPhysics
        for i, n in enumerate(self.items):
            self.park(n, i)
            prim = self.stage.GetPrimAtPath(f"{ITEMS_ROOT}/{n}")
            if prim.IsValid() and prim.HasAPI(UsdPhysics.RigidBodyAPI):
                UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(True)

    def _thaw(self, name):
        """let a parked item fall under gravity again, just before it is released"""
        from pxr import UsdPhysics
        prim = self.stage.GetPrimAtPath(f"{ITEMS_ROOT}/{name}")
        if prim.IsValid() and prim.HasAPI(UsdPhysics.RigidBodyAPI):
            UsdPhysics.RigidBodyAPI(prim).CreateKinematicEnabledAttr().Set(False)

    def release(self, name, y=0.0):
        self._thaw(name)
        self.place(name, (C.SPAWN_X, y, C.BELT_Z + 0.005))
        self.place(name, (C.SPAWN_X, y, C.BELT_Z + self._underside_gap(name) + 0.008))

    # -- laser gate ---------------------------------------------------------
    def laser(self):
        """name of whatever breaks the beam, else None. A real raycast, not a coordinate test."""
        hit = self._query.raycast_closest(
            [C.GATE_X, C.BEAM_Y0, C.BEAM_Z], [0.0, 1.0, 0.0], C.BEAM_Y1 - C.BEAM_Y0)
        if not hit or not hit.get("hit"):
            return None
        path = str(hit.get("rigidBody") or hit.get("collision") or "")
        for n in self.items:
            if f"{ITEMS_ROOT}/{n}" in path:
                return n
        return None

    def blade_path_busy(self, exclude):
        for n in self.items:
            if n == exclude:
                continue
            p = self.pose(n)
            if (C.BLADE_X0 - 0.12 < p[0] < C.BLADE_X1 + 0.12) and (-0.30 < p[1] < 0.45):
                return n
        return None

    async def divert(self, app, name, speed=None, max_wait_s=1.5):
        """full push cycle with both retract interlocks"""
        dt = 1.0 / 60.0
        t = await self.stroke(app, out=True, speed=speed)
        for _ in range(int(max_wait_s / dt)):                    # item off the main line
            if self.pose(name)[1] > 0.50:
                break
            await app.update_app_async(steps=1)
            t += dt
        held = 0.0
        for _ in range(int(max_wait_s / dt)):                    # path clear for the return
            if self.blade_path_busy(name) is None:
                break
            await app.update_app_async(steps=1)
            t += dt
            held += dt
        t += await self.stroke(app, out=False, speed=speed)
        return t, held

    # -- outcome ------------------------------------------------------------
    def where(self, name):
        """bin / branch / line-end / line, from the simulated pose"""
        p = self.pose(name)
        if C.BIN_X0 < p[0] < C.BIN_X1 and C.BIN_Y0 < p[1] < C.BIN_Y1 and p[2] < C.BIN_LIP_Z:
            return "bin"
        if p[2] < C.BELT_Z - 0.35:          # dropped off the end of the run
            return "line-end"
        if p[1] > 0.5:
            return "branch"
        if p[0] < C.MAIN_X0 + 0.35:         # the run stops at MAIN_X0, not beyond it
            return "line-end"
        return "line"