isaac/robozon_sorter/sim/spawner.py

"""Self-running item feeder: press Play and goods appear on the infeed belt one at a time,
spaced by a fixed pitch along the belt.

It hooks a PhysX step callback rather than living in an outer async loop, so the scene runs
on its own from the Play button - no driver script has to be babysitting it. The same
callback also drives the laser gate and the pusher when `route` is enabled.

Pitch is measured along the belt between consecutive items, so the release condition is
simply "the last one released has travelled PITCH from the spawn point".
"""
from __future__ import annotations

from .. import config as C


class AutoFeeder:
    def __init__(self, cell, order=None, pitch=None, loop=False,
                 route=None, on_event=None):
        """
        cell     : mechanics.Cell
        order    : release order; defaults to every loaded item
        pitch    : metres between consecutive items along the belt
        route    : dict name -> class; when given, class D is diverted by the pusher
        on_event : optional callback(kind, name, payload) for logging
        """
        self.cell = cell
        self.order = list(order or cell.items)
        self.pitch = pitch if pitch is not None else C.RELEASE_GAP
        self.loop = loop
        self.route = route or {}
        self.on_event = on_event
        self._sub = None
        self.reset()

    def reset(self):
        self.next_index = 0
        self.active = []
        self.released = []
        self.diverted = set()
        self.finished = {}
        self._busy = False          # a push cycle owns the blade until it completes
        self._cycle = None

    # ------------------------------------------------------------------ install
    def install(self):
        """subscribe to the physics step; from here on the cell runs itself on Play"""
        from omni.physx import get_physx_interface
        if self._sub is None:
            self._sub = get_physx_interface().subscribe_physics_step_events(self._on_step)
        return self

    def remove(self):
        self._sub = None

    def _emit(self, kind, name, payload=None):
        if self.on_event:
            self.on_event(kind, name, payload or {})

    # ------------------------------------------------------------------ per step
    def _on_step(self, dt):
        try:
            self._release_due()
            self._service_gate(dt)
            self._retire()
        except Exception as exc:                     # never let a callback kill the sim
            self._emit("error", "", {"exc": repr(exc)})

    def _release_due(self):
        if self._busy or self.next_index >= len(self.order):
            if self.loop and self.next_index >= len(self.order) and not self.active:
                self.next_index = 0
            return
        if self.active:
            travelled = C.SPAWN_X - float(self.cell.pose(self.active[-1])[0])
            if travelled < self.pitch:
                return
        name = self.order[self.next_index]
        self.cell.release(name)
        self.active.append(name)
        self.released.append(name)
        self.next_index += 1
        self._emit("release", name, {"pitch": self.pitch})

    def _service_gate(self, dt):
        """laser gate -> pusher, as a small state machine so it spans several steps"""
        if self._cycle is not None:
            self._step_cycle(dt)
            return
        for name in list(self.active):
            if name in self.diverted or self.route.get(name) != "D":
                continue
            if self.cell.laser() == name:
                self._cycle = dict(name=name, phase="extend", t=0.0,
                                   y=C.BLADE_HOME_Y, held=0.0)
                self._busy = True
                self._emit("gate", name, {})
                return

    def _step_cycle(self, dt):
        c = self._cycle
        name = c["name"]
        speed = C.PUSHER_SPEED
        if c["phase"] == "extend":
            c["y"] = min(C.BLADE_OUT_Y, c["y"] + speed * dt)
            self.cell.blade_to(c["y"])
            if c["y"] >= C.BLADE_OUT_Y - 1e-6:
                c["phase"] = "clear"
        elif c["phase"] == "clear":
            c["t"] += dt
            if float(self.cell.pose(name)[1]) > 0.50 or c["t"] > 1.5:
                c["phase"] = "wait"
                c["t"] = 0.0
        elif c["phase"] == "wait":
            # do not sweep the blade back through whatever has already arrived
            busy = self.cell.blade_path_busy(name)
            c["t"] += dt
            if busy is None or c["t"] > 1.5:
                c["held"] = c["t"]
                c["phase"] = "retract"
        elif c["phase"] == "retract":
            c["y"] = max(C.BLADE_HOME_Y, c["y"] - speed * dt)
            self.cell.blade_to(c["y"])
            if c["y"] <= C.BLADE_HOME_Y + 1e-6:
                self.diverted.add(name)
                self._busy = False
                self._cycle = None
                self._emit("divert", name, {"held": round(c["held"], 3)})

    def _retire(self):
        for name in list(self.active):
            place = self.cell.where(name)
            if place in ("bin", "line-end"):
                self.finished[name] = place
                self.active.remove(name)
                self._emit("done", name, {"where": place})