source/bimanual/yam/conditions.py

"""Reusable success conditions for the YAM task suite.

Modelled on RoboLab's ``robolab/core/task/conditionals.py``: a task declares *what* counts as
success as a list of small predicates, instead of open-coding the geometry test in the demo
script. Each predicate takes the world state dict produced by a task script and returns a bool,
so the same condition can be reused by the runner, by an eval harness, or by a policy's reward.

State dict shape (env-local metres, table top at ``TABLE_TOP``)::

    {"objects": {"grape": np.array([x, y, z]), ...},
     "quats":   {"pot": np.array([w, x, y, z]), ...},     # optional
     "regions": {"basket": {"xy": (x, y), "half": 0.085, "top_z": 0.52}},
     "start_z": {"grape": 0.449}}
"""
from __future__ import annotations

from dataclasses import dataclass
from typing import Callable

import numpy as np

TABLE_TOP = 0.45


def _pos(state, name):
    p = state.get("objects", {}).get(name)
    if p is None:
        raise KeyError(f"object {name!r} not in state; have {sorted(state.get('objects', {}))}")
    return np.asarray(p, dtype=float)


def _region(state, name):
    r = state.get("regions", {}).get(name)
    if r is None:
        raise KeyError(f"region {name!r} not in state; have {sorted(state.get('regions', {}))}")
    return r


def object_in_region(object: str, region: str, pad: float = 0.02, below_top: float = 0.16):
    """Object's centre is inside the region footprint and low enough to be *in* it, not on top."""
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        half = float(r["half"])+pad
        dx, dy = abs(p[0]-r["xy"][0]), abs(p[1]-r["xy"][1])
        ok_xy = dx < half and dy < half
        ok_z = p[2] < TABLE_TOP+below_top
        # Say WHICH test failed. "not in the region" is ambiguous between missed sideways and
        # left sitting on the rim, and those want opposite fixes.
        if not (ok_xy and ok_z):
            print(f"[cond] {object} at ({p[0]:+.3f},{p[1]:+.3f},{p[2]:.3f}) vs {region} "
                  f"centre ({r['xy'][0]:+.3f},{r['xy'][1]:+.3f}) half={half:.3f}: "
                  f"dx={dx:.3f} dy={dy:.3f} {'' if ok_xy else 'OUTSIDE-XY '}"
                  f"{'' if ok_z else 'TOO-HIGH(sitting on the rim?)'}", flush=True)
        return ok_xy and ok_z
    return cond


def object_on_surface(object: str, region: str, pad: float = 0.03, tol: float = 0.01):
    """Object rests ON the region's top surface (a scale, plate, coaster)."""
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        half = float(r["half"])+pad
        return (abs(p[0]-r["xy"][0]) < half and abs(p[1]-r["xy"][1]) < half
                and p[2] > float(r["top_z"])-tol)
    return cond


def object_not_in_region(object: str, region: str, pad: float = 0.02):
    """A distractor must be left alone -- the negative half of a semantic-selection task."""
    inner = object_in_region(object, region, pad=pad)
    return lambda state: not inner(state)


def object_lifted(object: str, min_dz: float = 0.05):
    """Object rose by at least ``min_dz`` from where it started (a real grasp, not a nudge)."""
    def cond(state):
        z0 = state.get("start_z", {}).get(object)
        peak = state.get("peak_z", {}).get(object, _pos(state, object)[2])
        return z0 is not None and (float(peak)-float(z0)) > min_dz
    return cond


def object_seated_in_hole(object: str, region: str, xy_tol: float = 0.02, half_h: float = 0.065):
    """Peg inserted: centred on the hole AND down at table level, not perched on the socket."""
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        return (abs(p[0]-r["xy"][0]) < xy_tol and abs(p[1]-r["xy"][1]) < xy_tol
                and p[2] < TABLE_TOP+half_h+0.02)
    return cond


def objects_stacked(lower: str, upper: str, min_dz: float = 0.03, xy_tol: float = 0.03):
    """``upper`` sits on top of ``lower``, roughly aligned."""
    def cond(state):
        a = _pos(state, lower); b = _pos(state, upper)
        return (b[2]-a[2]) > min_dz and abs(b[0]-a[0]) < xy_tol and abs(b[1]-a[1]) < xy_tol
    return cond


def _up_axis(q):
    w, x, y, z = [float(v) for v in q]
    return np.array([2*(x*z+y*w), 2*(y*z-x*w), 1-2*(x*x+y*y)])


def object_level(object: str, max_tilt_deg: float = 25.0):
    """Object has not tilted away from HOW IT WAS RESTING.

    Not "its local +z points up": a mesh that needs a 90 deg roll to stand upright (most RoboTwin
    assets) has its local +z horizontal by construction, so comparing to world up marks a
    perfectly level pot as 90 deg tilted.
    """
    def cond(state):
        q = state.get("quats", {}).get(object)
        q0 = state.get("start_quats", {}).get(object)
        if q is None:
            return False
        now = _up_axis(q)
        ref = _up_axis(q0) if q0 is not None else np.array([0.0, 0.0, 1.0])
        cosang = float(np.dot(now, ref)/((np.linalg.norm(now)*np.linalg.norm(ref))+1e-9))
        return float(np.degrees(np.arccos(np.clip(cosang, -1.0, 1.0)))) < max_tilt_deg
    return cond


def count_in_region(objects: list[str], region: str, at_least: int, pad: float = 0.03,
                    below_top: float = 0.10):
    """At least N of the listed objects ended up in the region (beads swept/poured).

    `below_top` is not optional in practice: without a height bound this counts a bead that is
    still sitting in the vessel being carried over the target, which reports a successful pour
    for an episode that never poured.
    """
    def cond(state):
        r = _region(state, region); half = float(r["half"])+pad
        top = float(r.get("top_z", TABLE_TOP))
        n, where = 0, []
        for o in objects:
            p = state.get("objects", {}).get(o)
            if p is None:
                where.append(f"{o}=missing")
                continue
            p = np.asarray(p, dtype=float)
            ok = (abs(p[0]-r["xy"][0]) < half and abs(p[1]-r["xy"][1]) < half
                  and p[2] < top+below_top)
            where.append(f"{o}=({p[0]:+.3f},{p[1]:+.3f},{p[2]:.3f}){'IN' if ok else ''}")
            n += int(ok)
        print(f"[cond] {region} half={half:.3f} top={top:.3f} -> {n}/{len(objects)} in: "
              + " ".join(where), flush=True)
        return n >= at_least
    return cond


def object_moved_to(object: str, region: str, tol: float = 0.07):
    """Object ended near the target -- used by push/slide, where nothing is grasped."""
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        return float(np.hypot(p[0]-r["xy"][0], p[1]-r["xy"][1])) < tol
    return cond


def joint_moved(fixture: str, index: int = 0, by: float = 0.5):
    """A jointed fixture's joint travelled at least `by` from where the episode started.

    Radians for a hinge (door, laptop lid, toggle), metres for a slider (drawer). Signed
    magnitude, so it reads the same whether the task is to open or to close.
    """
    def cond(state):
        q0 = state.get("start_joints", {}).get((fixture, index))
        q = state.get("joints", {}).get((fixture, index))
        if q is None or q0 is None:
            return False
        return abs(float(q)-float(q0)) >= float(by)
    return cond


def link_span_at_least(fixture: str, a: int, b: int, min_span: float):
    """Two links of a chain ended at least `min_span` apart -- i.e. the rope is taut.

    Measured on the rope's own geometry rather than on where the hands are: a hand can be in the
    right place while the rope hangs slack between them.
    """
    def cond(state):
        pa = state.get("links", {}).get((fixture, a))
        pb = state.get("links", {}).get((fixture, b))
        if pa is None or pb is None:
            print(f"[cond] {fixture}: link {a}/{b} not in state", flush=True); return False
        d = float(np.linalg.norm(np.asarray(pa)-np.asarray(pb)))
        print(f"[cond] {fixture} span link{a}-link{b} = {d:.3f} m (need {min_span:.3f})", flush=True)
        return d >= float(min_span)
    return cond


def link_above(fixture: str, index: int, min_z: float = 0.04):
    """One link of a chain is at least `min_z` above the table."""
    def cond(state):
        p = state.get("links", {}).get((fixture, index))
        ok = p is not None and float(p[2]) >= TABLE_TOP+float(min_z)
        print(f"[cond] {fixture} link{index} z="
              f"{'?' if p is None else round(float(p[2])-TABLE_TOP,3)} m above the table "
              f"(need {min_z})", flush=True)
        return ok
    return cond


def joint_driven(fixture: str, index: int = 0, by: float = 0.3):
    """The ROBOT moved this joint by `by`, measured across the manipulation itself.

    joint_moved() compares against a baseline the env recorded earlier, which also counts the
    fixture settling under gravity and any pre-action blow-up. This reads the delta the solver
    measured between "about to act" and "finished acting", and refuses a joint that ended
    outside its own limits.
    """
    def cond(state):
        d = state.get("joint_delta", {}).get((fixture, index))
        sane = state.get("joint_sane", {}).get(fixture, True)
        print(f"[cond] {fixture} joint {index}: robot moved it "
              f"{'none' if d is None else round(float(d), 3)} (need {by}), "
              f"within limits={sane}", flush=True)
        return d is not None and float(d) >= float(by) and sane
    return cond


def joint_at_least(fixture: str, index: int = 0, value: float = 0.6):
    """A jointed fixture reached an absolute joint value (a door actually standing open)."""
    def cond(state):
        q = state.get("joints", {}).get((fixture, index))
        return q is not None and float(q) >= float(value)
    return cond


def object_tipped_over(object: str, region: str, min_tilt_deg: float = 75.0, xy_tol: float = 0.12):
    """Vessel was brought over a target and rotated far enough to pour out of.

    This is the honest success test for pouring in a rigid-body sim. A vessel that has to be
    LIFTED must be a convex dynamic body -- PhysX will not simulate a hollow carried mesh -- so
    its interior is solid and anything "inside" is pushed out to the table before the episode
    starts. What can be checked is the pour itself: the vessel above the target, tipped past the
    angle at which its contents would leave it.
    """
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        q = state.get("peak_tilt", {}).get(object)
        d = float(np.hypot(p[0]-r["xy"][0], p[1]-r["xy"][1]))
        ok = q is not None and float(q) >= float(min_tilt_deg) and d < float(xy_tol)
        print(f"[cond] {object}: peak tilt {q if q is None else round(float(q),1)} deg "
              f"(need {min_tilt_deg}), ended {d:.3f} m from {region} (need <{xy_tol})", flush=True)
        return ok
    return cond


def object_settled_on(object: str, region: str, max_dz: float = 0.03):
    """Object came to rest ON a fixture's base rather than hanging up partway.

    A ring balanced on the tip of a post satisfies "centred over the post" just as well as one
    that slid all the way down, so the height above the base is what separates them.
    """
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        return p[2]-float(r["top_z"]) < float(max_dz)
    return cond


def object_clear_of(object: str, region: str, extent: float, margin: float = 0.0):
    """Object has been drawn ENTIRELY out of a fixture -- the inverse of object_moved_to.

    `extent` is the object's own half-length along the pull axis, so the check is "its trailing
    end is past the fixture's mouth", not merely "its centre moved a bit".
    """
    def cond(state):
        p = _pos(state, object); r = _region(state, region)
        d = float(np.hypot(p[0]-r["xy"][0], p[1]-r["xy"][1]))
        return d > float(r.get("half", 0.0))+float(extent)-float(margin)
    return cond


@dataclass
class Subtask:
    """A named step whose ``conditions`` must ALL hold (RoboLab's Subtask, trimmed)."""
    name: str
    conditions: list[Callable]

    def satisfied(self, state) -> bool:
        return all(c(state) for c in self.conditions)


def evaluate(subtasks: list[Subtask], state) -> tuple[bool, dict[str, bool]]:
    """Returns (all_passed, per-subtask results)."""
    per = {s.name: s.satisfied(state) for s in subtasks}
    return all(per.values()), per


def labelled(label: str, cond):
    """Attach a human-readable label to a predicate, for the per-condition PASS/FAIL report."""
    cond.label = label
    return cond