isaac/robozon_sorter/sim/plow_contact.py

"""Contact sensing on the plow blade.

**Why a contact report and not another beam.** The cell already has two through-beams: the
laser gate before the pusher and the arming beam at x = -6.30 that pre-positions the plow.
Both answer "something is about to arrive". Neither can answer "the blade is now touching
*this* item", and that is the question that matters at the plow, because the arm is only
useful while it is actually in contact - before that it is waving at nothing, and after it
the item is already committed to a lane. A beam at the blade would also be broken by the
blade itself as it swings, which is the trap the gate beam at y = -0.24 was placed to dodge.

So the sensor is a **PhysX contact report on the arm body**
(``PhysxSchema.PhysxContactReportAPI``). It fires on the real collision pair, names both
bodies, and needs no extra geometry that could foul the belt. Isaac's
``sensors.experimental.physics.Contact`` wraps the same mechanism with an authored prim and
a threshold; the raw report is used here because the plow needs the *identity* of what it
touched, which is what carries the class through.

**Keeping the class.** Classification happens once, far upstream under the camera portal.
That verdict is stored per item and travels with it:

    camera portal ──▶ classes[item] = "B" | "C" | "D"
                          │
    arming beam  ────────▶ pre-position the blade for that class
                          │
    blade contact ───────▶ CONFIRM against the same stored class, and hold the side while
                           contact lasts - the item is steered by the class it was given,
                           not by anything re-derived at the blade

:class:`PlowContact` therefore takes the same ``classes`` mapping the sorter uses, and
reports, per touch: which item, what class it carries, the blade angle at first touch, and
how long contact lasted. A touch whose class is unknown is reported as such rather than
guessed - an unclassified item must not be steered anywhere.
"""
from __future__ import annotations

from pxr import PhysicsSchemaTools, PhysxSchema

from .. import config as C

ITEMS_PREFIX = "/World/Items/"


class PlowContact:
    """PhysX contact reporting on the plow arm, resolved to item + class"""

    def __init__(self, stage, classes: dict, arm_path: str | None = None,
                 plow=None, threshold: float = 0.0):
        """
        classes : the SAME dict the sorter steers by - vision writes into it, so the
                  sensor sees whatever verdict the item is carrying at the moment of touch
        plow    : optional sim.plow.Plow, so the angle at contact can be recorded
        """
        self.stage = stage
        self.classes = classes
        self.plow = plow
        self.arm_path = arm_path or C.PLOW_ARM

        prim = stage.GetPrimAtPath(self.arm_path)
        if not prim.IsValid():
            raise RuntimeError(f"{self.arm_path} missing - is this plow_cell.usd?")
        api = PhysxSchema.PhysxContactReportAPI.Apply(prim)
        api.CreateThresholdAttr().Set(float(threshold))     # 0 = report every touch

        self.touches: dict[str, dict] = {}      # item -> first/last touch record
        self.in_contact: set[str] = set()
        self.events: list[dict] = []
        self._t = 0.0
        self._sub = None

    # -- lifecycle ----------------------------------------------------------
    def install(self):
        from omni.physx import get_physx_simulation_interface
        if self._sub is None:
            self._sub = get_physx_simulation_interface(
            ).subscribe_contact_report_events(self._on_report)
        return self

    def remove(self):
        self._sub = None

    def tick(self, dt):
        """advance the sensor's clock; contact reports carry no timestamp of their own"""
        self._t += dt

    # -- the report ---------------------------------------------------------
    def _item_of(self, path: str):
        if not path.startswith(ITEMS_PREFIX):
            return None
        name = path[len(ITEMS_PREFIX):].split("/")[0]
        return name or None

    def _on_report(self, contact_headers, contact_data):
        touching = set()
        for h in contact_headers:
            a0 = str(PhysicsSchemaTools.intToSdfPath(h.actor0))
            a1 = str(PhysicsSchemaTools.intToSdfPath(h.actor1))
            if self.arm_path not in (a0, a1):
                continue
            other = a1 if self.arm_path == a0 else a0
            name = self._item_of(other)
            if name is None:                    # the blade also brushes belts and rails
                continue
            touching.add(name)
            self._register(name)
        # contact that has ended
        for gone in self.in_contact - touching:
            rec = self.touches.get(gone)
            if rec is not None:
                rec["released_t"] = round(self._t, 3)
                rec["duration"] = round(self._t - rec["first_t"], 3)
        self.in_contact = touching

    def _register(self, name):
        cls = self.classes.get(name)
        angle = round(self.plow.angle, 1) if self.plow is not None else None
        rec = self.touches.get(name)
        if rec is None:
            rec = dict(item=name, cls=cls, classified=cls is not None,
                       first_t=round(self._t, 3), angle_at_touch=angle,
                       commanded_at_touch=(round(self.plow.commanded, 1)
                                           if self.plow is not None else None),
                       angle_min=angle, angle_max=angle,
                       released_t=None, duration=None, samples=0)
            self.touches[name] = rec
            self.events.append(dict(t=rec["first_t"], item=name, cls=cls,
                                    angle=angle, kind="touch"))
        rec["samples"] += 1
        rec["cls"] = cls if cls is not None else rec["cls"]
        if angle is not None:
            rec["angle_min"] = min(rec["angle_min"], angle)
            rec["angle_max"] = max(rec["angle_max"], angle)

    # -- what the plow asks it ----------------------------------------------
    def is_touching(self, name: str) -> bool:
        return name in self.in_contact

    def touched(self, name: str) -> bool:
        return name in self.touches

    def side_for(self, name: str, mapping: dict, swing: float):
        """the angle this item's stored class asks for, or None if it has no class.

        Deliberately returns None rather than 0 for an unknown class: 0 is a real command
        (drive straight on) and must not double as "no idea".
        """
        cls = self.classes.get(name)
        if cls is None:
            return None
        want = mapping.get(cls, "straight")
        return {"pos": swing, "neg": -swing}.get(want, 0.0)

    def report(self):
        return dict(touches=list(self.touches.values()), events=self.events)