isaac/robozon_sorter/sim/plow.py

"""The plow diverter (``DiverterEnd``) - the blade at the far end of the main run.

Mechanically it is the opposite of the pusher. The pusher is a kinematic slab shoved
across the belt from script; the plow is a **dynamic arm on a revolute joint driven by an
angular force drive**, so it is compliant - it yields on contact instead of teleporting
through cargo. That is why this module commands a drive target rather than writing a
transform the way ``mechanics.Cell.blade_to`` does.

Two consequences of that choice, both load-bearing:

1. **A drive target is a request, not a position.** The arm arrives when the solver gets it
   there, and USD drive writes reach PhysX with a lag (about a second in the worst case
   measured on the pusher's prismatic joint). Never assume the blade is where you last
   commanded it - read :meth:`Plow.angle`, which measures the arm's actual pose.

2. **Rate is not free.** The authored graph swings 30 deg in 7 ms (72 rad/s). At that rate
   the blade is an impulse and throws goods off the line. :meth:`Plow.step_toward` ramps the
   target at ``config.PLOW_RATE`` instead, which is what makes the motion sortable.

The scene keeps its authored ``DiverterAnimGraph``, so pressing Play alone makes the cell
demonstrate itself. ``plow_cell.prepare(..., script_control=True)`` switches that graph off;
until it is off, the graph rewrites the drive target every tick and fights this module.
"""
from __future__ import annotations

import math

from pxr import UsdGeom, UsdPhysics

from .. import config as C


def _yaw_deg(quat) -> float:
    """Z rotation of a [w, x, y, z] quaternion, in degrees"""
    w, x, y, z = (float(v) for v in quat)
    return math.degrees(math.atan2(2.0 * (w * z + x * y), 1.0 - 2.0 * (y * y + z * z)))


class Plow:
    """Angular control of the plow arm.

    ``angle`` is measured from the arm's rest pose, so it is signed the same way as the
    authored throw: positive swings one way into the lane, negative the other.
    """

    def __init__(self, stage, hinge_path: str | None = None, arm_path: str | None = None,
                 base_path: str | None = None, kinematic: bool = True):
        self.stage = stage
        self.hinge_path = hinge_path or C.PLOW_HINGE
        self.arm_path = arm_path or C.PLOW_ARM

        hinge = stage.GetPrimAtPath(self.hinge_path)
        if not hinge.IsValid():
            raise RuntimeError(
                f"{self.hinge_path} missing - plow_cell.usd is the scene with the plow; "
                "sorter.usd carries the pusher only")
        self.drive = UsdPhysics.DriveAPI(hinge, "angular")
        if not self.drive:
            raise RuntimeError(f"{self.hinge_path} has no angular drive to command")

        from isaacsim.core.experimental.prims import RigidPrim
        self._arm = RigidPrim(paths=[self.arm_path])
        # Reference the angle to the *base*, not to whatever pose the arm happened to hold
        # when this object was built. The base is kinematic and both bodies read yaw +180
        # at rest, so `yaw(arm) - yaw(base)` is the true joint angle and reads 0 at rest.
        #
        # Taking a snapshot instead was wrong and hid every other plow fault: a run that
        # started with the arm left at -30 from the previous run reported "commanded 0.0 ->
        # reached +30.47" and "commanded +30.0 -> reached -34.56", which looks like a
        # broken drive rather than a broken measurement.
        self._base = RigidPrim(paths=[base_path or C.PLOW_BASE])
        self.commanded = 0.0

        # Kinematic mode: rotate the arm directly instead of asking a force drive to hold
        # an angle. The compliant drive was tuned three times (120000 / 3000 / 300) and
        # never held its target - at 3000 it rang between +-21.4 deg at 99 deg/s, faster
        # than the 76.4 deg/s ramp commanding it, which is the drive moving the arm rather
        # than the command. A kinematic arm turned at the ramp rate goes exactly where it
        # is put, which is what the pusher blade has always done.
        #
        # It gives up compliance, so the arm no longer yields on contact. That is safe here
        # only because the tip speed matches the belt (0.8 m/s): the blade leans goods over
        # at their own speed rather than batting them. MAX_DEPENETRATION still caps how
        # violently PhysX may separate a deep overlap.
        self._rot_op = None
        if kinematic:
            self._rot_op = self._ensure_rot_op()

    def _ensure_rot_op(self):
        """the arm's own rotateZ op, created if the authored prim has none.

        The hinge sits at the arm's origin (localPos0 = localPos1 = 0), so turning the arm
        about its own Z reproduces the joint exactly.
        """
        xf = UsdGeom.Xformable(self.stage.GetPrimAtPath(self.arm_path))
        for op in xf.GetOrderedXformOps():
            if op.GetOpType() == UsdGeom.XformOp.TypeRotateZ:
                return op
        return xf.AddRotateZOp()

    # -- state --------------------------------------------------------------
    def _yaw_of(self, prim) -> float:
        """world yaw from the SIMULATED pose (never XformCache: during simulation that
        returns the authored transform and the arm looks frozen)"""
        return _yaw_deg(prim.get_world_poses()[1].numpy()[0])

    @property
    def angle(self) -> float:
        """true joint angle in degrees: the arm's yaw relative to the base it hinges on"""
        d = self._yaw_of(self._arm) - self._yaw_of(self._base)
        return (d + 180.0) % 360.0 - 180.0

    def at(self, deg: float, tol: float = 1.0) -> bool:
        return abs(self.angle - deg) <= tol

    # -- command ------------------------------------------------------------
    def target(self, deg: float, velocity: float | None = None) -> float:
        """command the drive; the value is clamped inside the joint's own limit"""
        deg = max(-C.PLOW_LIMIT, min(C.PLOW_LIMIT, float(deg)))
        if self._rot_op is not None:            # kinematic: put the arm there
            self._rot_op.Set(float(deg))
        else:                                   # compliant: ask the drive to get there
            self.drive.GetTargetPositionAttr().Set(deg)
            if velocity is not None:
                self.drive.GetTargetVelocityAttr().Set(float(velocity))
        self.commanded = deg
        return deg

    def home(self) -> float:
        """centre the blade, out of the lane"""
        return self.target(0.0, velocity=0.0)

    def gains(self, stiffness=None, damping=None, max_force=None):
        """re-apply the authored gains, or override them for an experiment"""
        self.drive.GetStiffnessAttr().Set(float(
            C.PLOW_STIFFNESS if stiffness is None else stiffness))
        self.drive.GetDampingAttr().Set(float(
            C.PLOW_DAMPING if damping is None else damping))
        self.drive.GetMaxForceAttr().Set(float(
            C.PLOW_MAX_FORCE if max_force is None else max_force))

    # -- motion -------------------------------------------------------------
    def step_toward(self, deg: float, dt: float, rate: float | None = None) -> bool:
        """advance the commanded target one physics step toward `deg`.

        Ramping the target is what keeps the blade sortable: commanding the endpoint
        outright makes the solver deliver it as an impulse. Returns True once the command
        has reached `deg` - the arm itself follows a little later, so gate on
        :meth:`at` if you need the blade physically there.
        """
        rate = C.PLOW_RATE if rate is None else rate
        step = rate * dt
        delta = deg - self.commanded
        if abs(delta) <= step:
            self.target(deg, velocity=0.0)
            return True
        self.target(self.commanded + math.copysign(step, delta),
                    velocity=math.copysign(rate, delta))
        return False

    async def swing(self, app, deg: float, rate: float | None = None,
                    settle_s: float = 0.5, dt: float = 1.0 / 60.0) -> float:
        """drive the blade to `deg` and wait for the arm to actually arrive"""
        rate = C.PLOW_RATE if rate is None else rate
        t = 0.0
        while not self.step_toward(deg, dt, rate):
            await app.update_app_async(steps=1)
            t += dt
        for _ in range(int(settle_s / dt)):          # the arm lags the command
            if self.at(deg):
                break
            await app.update_app_async(steps=1)
            t += dt
        return t

    async def divert(self, app, side: float = 1.0, dwell_s: float | None = None,
                     rate: float | None = None) -> float:
        """full cycle: swing into the lane, hold, return to centre"""
        dwell_s = C.PLOW_HOLD if dwell_s is None else dwell_s
        deg = math.copysign(C.PLOW_SWING, side)
        t = await self.swing(app, deg, rate)
        for _ in range(int(dwell_s / (1.0 / 60.0))):
            await app.update_app_async(steps=1)
            t += 1.0 / 60.0
        t += await self.swing(app, 0.0, rate)
        return t

    # -- the authored profile, in Python -------------------------------------
    @staticmethod
    def authored_profile(t: float, rate: float | None = None) -> tuple[float, float]:
        """(target_deg, target_deg_per_s) of the scene's own OmniGraph loop at time `t`.

        Reimplemented so the demo motion is available from code. `rate` defaults to
        ``config.PLOW_RATE`` rather than the authored 4125 deg/s; pass
        ``config.PLOW_RATE_AUTHORED`` to reproduce the scene exactly, impulse and all.
        """
        rate = C.PLOW_RATE if rate is None else rate
        a = C.PLOW_SWING
        hold, ts = C.PLOW_HOLD, a / rate
        segs = [(hold, 0.0, 0.0), (ts, a, rate), (hold, a, 0.0), (ts, 0.0, -rate),
                (hold, 0.0, 0.0), (ts, -a, -rate), (hold, -a, 0.0), (ts, 0.0, rate),
                (hold, 0.0, 0.0)]
        period = sum(s[0] for s in segs)
        c, t0, pos0 = t % period, 0.0, 0.0
        for dur, pos1, vel in segs:
            if c <= t0 + dur + 1e-9:
                u = 0.0 if dur <= 0 else (c - t0) / dur
                pos = pos0 + (pos1 - pos0) * u if abs(vel) > 1e-6 else pos1
                return float(pos), float(vel)
            t0, pos0 = t0 + dur, pos1
        return 0.0, 0.0