src/evaluation.py - master - koda/FiVI

src/evaluation.py
454 строки17 KB
koda
Prepare production-ready grasp planner
18 июл 2026, 09:39
18 июл 2026, 09:39f540458
from __future__ import annotations

from dataclasses import dataclass
import mujoco  # type: ignore[import-untyped]
import numpy as np

from .ik import FINGER_IDS, FINGERTIP_GEOMS, FloatArray, expand_coupled_joints
from .planner import (
    Failure,
    Plan,
    _adaptive_path_alphas,
    _kinematically_adjacent,
)
from .runtime import Runtime
from .scene import Scene


EXACT_CLEARANCE_TOLERANCE_M = 1.0e-6
EXACT_CONTACT_TOLERANCE_M = 0.007
EXACT_FINGERTIP_GAP_M = 0.010
EXACT_MINIMUM_PATH_SAMPLES = 12
EXACT_PATH_SAMPLES = EXACT_MINIMUM_PATH_SAMPLES
EXACT_MAX_PATH_JOINT_STEP_RAD = 0.0125
SELF_COLLISION_TOLERANCE_M = 2.0e-4
_GEOM_DISTANCE_LIMIT_M = 1.0
_STATE_SPEC = int(mujoco.mjtState.mjSTATE_INTEGRATION)


@dataclass(frozen=True, slots=True)
class Evaluation:
    contact_errors_m: FloatArray
    fingertip_clearances_m: FloatArray
    true_clearance_m: float
    object_pose_unchanged: bool
    path_clearance_m: float = np.inf
    self_clearance_m: float = np.inf
    path_self_clearance_m: float = np.inf
    geometry_supported: bool = True

    def __post_init__(self) -> None:
        contact_errors = np.array(self.contact_errors_m, dtype=np.float64, copy=True)
        fingertip_clearances = np.array(
            self.fingertip_clearances_m,
            dtype=np.float64,
            copy=True,
        )
        if contact_errors.ndim != 1 or not bool(np.all(np.isfinite(contact_errors))):
            raise ValueError("contact_errors_m must be a finite vector")
        if fingertip_clearances.ndim != 1 or not bool(
            np.all(np.isfinite(fingertip_clearances))
        ):
            raise ValueError("fingertip_clearances_m must be a finite vector")
        scalars = np.asarray(
            (
                self.true_clearance_m,
                self.path_clearance_m,
                self.self_clearance_m,
                self.path_self_clearance_m,
            ),
            dtype=np.float64,
        )
        if bool(np.any(np.isnan(scalars))) or bool(np.any(np.isneginf(scalars))):
            raise ValueError("evaluation clearances must be finite or positive infinity")
        contact_errors.flags.writeable = False
        fingertip_clearances.flags.writeable = False
        object.__setattr__(self, "contact_errors_m", contact_errors)
        object.__setattr__(self, "fingertip_clearances_m", fingertip_clearances)

    @property
    def accepted(self) -> bool:
        clearances = np.asarray(
            (
                self.true_clearance_m,
                self.path_clearance_m,
                self.self_clearance_m,
                self.path_self_clearance_m,
            ),
            dtype=np.float64,
        )
        if (
            not self.geometry_supported
            or not self.object_pose_unchanged
            or self.contact_errors_m.size == 0
            or self.fingertip_clearances_m.size == 0
            or not bool(np.all(np.isfinite(clearances)))
        ):
            return False
        return bool(
            float(np.max(self.contact_errors_m)) <= EXACT_CONTACT_TOLERANCE_M
            and float(np.min(self.fingertip_clearances_m))
            >= -EXACT_CLEARANCE_TOLERANCE_M
            and float(np.max(self.fingertip_clearances_m)) <= EXACT_FINGERTIP_GAP_M
            and self.true_clearance_m >= -EXACT_CLEARANCE_TOLERANCE_M
            and self.path_clearance_m >= -EXACT_CLEARANCE_TOLERANCE_M
            and self.self_clearance_m >= -SELF_COLLISION_TOLERANCE_M
            and self.path_self_clearance_m >= -SELF_COLLISION_TOLERANCE_M
        )

    @property
    def failure(self) -> Failure | None:
        return None if self.accepted else Failure.POSTCHECK

    @property
    def sampled_path_clearance_m(self) -> float:
        return self.path_clearance_m

    @property
    def sampled_path_self_clearance_m(self) -> float:
        return self.path_self_clearance_m


def _segment_distances(
    point: FloatArray,
    starts: FloatArray,
    ends: FloatArray,
) -> FloatArray:
    directions = ends - starts
    denominator = np.einsum("ij,ij->i", directions, directions)
    fraction = np.divide(
        np.einsum("ij,ij->i", point - starts, directions),
        denominator,
        out=np.zeros_like(denominator),
        where=denominator > np.finfo(np.float64).eps,
    )
    fraction = np.clip(fraction, 0.0, 1.0)
    closest = starts + fraction[:, None] * directions
    distances: FloatArray = np.linalg.norm(closest - point, axis=1)
    return distances


def _mesh_distance(runtime: Runtime, geom_id: int, point_local: FloatArray) -> float:
    mesh_id = int(runtime.model.geom_dataid[geom_id])
    vertex_start = int(runtime.model.mesh_vertadr[mesh_id])
    vertex_count = int(runtime.model.mesh_vertnum[mesh_id])
    face_start = int(runtime.model.mesh_faceadr[mesh_id])
    face_count = int(runtime.model.mesh_facenum[mesh_id])
    vertices = np.asarray(
        runtime.model.mesh_vert[vertex_start : vertex_start + vertex_count],
        dtype=np.float64,
    )
    faces = np.asarray(
        runtime.model.mesh_face[face_start : face_start + face_count],
        dtype=np.intp,
    )
    if not face_count:
        return float(np.min(np.linalg.norm(vertices - point_local, axis=1)))
    triangles = vertices[faces]
    first = triangles[:, 0]
    second = triangles[:, 1]
    third = triangles[:, 2]
    first_edge = second - first
    second_edge = third - first
    relative = point_local - first
    dot00 = np.einsum("ij,ij->i", first_edge, first_edge)
    dot01 = np.einsum("ij,ij->i", first_edge, second_edge)
    dot11 = np.einsum("ij,ij->i", second_edge, second_edge)
    dot20 = np.einsum("ij,ij->i", relative, first_edge)
    dot21 = np.einsum("ij,ij->i", relative, second_edge)
    denominator = dot00 * dot11 - dot01 * dot01
    inverse = np.divide(
        1.0,
        denominator,
        out=np.zeros_like(denominator),
        where=np.abs(denominator) > np.finfo(np.float64).eps,
    )
    first_coordinate = (dot11 * dot20 - dot01 * dot21) * inverse
    second_coordinate = (dot00 * dot21 - dot01 * dot20) * inverse
    inside = (
        (first_coordinate >= 0.0)
        & (second_coordinate >= 0.0)
        & (first_coordinate + second_coordinate <= 1.0)
        & (np.abs(denominator) > np.finfo(np.float64).eps)
    )
    normals = np.cross(first_edge, second_edge)
    normal_length = np.linalg.norm(normals, axis=1)
    plane_distance = np.divide(
        np.abs(np.einsum("ij,ij->i", relative, normals)),
        normal_length,
        out=np.full_like(normal_length, np.inf),
        where=normal_length > np.finfo(np.float64).eps,
    )
    edge_distance = np.minimum(
        _segment_distances(point_local, first, second),
        np.minimum(
            _segment_distances(point_local, second, third),
            _segment_distances(point_local, third, first),
        ),
    )
    return float(np.min(np.where(inside, plane_distance, edge_distance)))


def _signed_distance(
    runtime: Runtime,
    geom_id: int,
    point: FloatArray,
) -> tuple[float, bool]:
    rotation = np.asarray(runtime.data.geom_xmat[geom_id], dtype=np.float64).reshape(3, 3)
    local = (point - runtime.data.geom_xpos[geom_id]) @ rotation
    size = np.asarray(runtime.model.geom_size[geom_id], dtype=np.float64)
    geom_type = int(runtime.model.geom_type[geom_id])
    if geom_type == int(mujoco.mjtGeom.mjGEOM_BOX):
        delta = np.abs(local) - size
        return (
            float(
                np.linalg.norm(np.maximum(delta, 0.0))
                + min(float(np.max(delta)), 0.0)
            ),
            True,
        )
    if geom_type == int(mujoco.mjtGeom.mjGEOM_SPHERE):
        return float(np.linalg.norm(local) - size[0]), True
    if geom_type == int(mujoco.mjtGeom.mjGEOM_CAPSULE):
        axial = max(abs(float(local[2])) - float(size[1]), 0.0)
        radial = float(np.linalg.norm(local[:2]))
        return float(np.hypot(radial, axial) - size[0]), True
    if geom_type == int(mujoco.mjtGeom.mjGEOM_CYLINDER):
        delta = np.array(
            (np.linalg.norm(local[:2]) - size[0], abs(local[2]) - size[1])
        )
        return (
            float(
                np.linalg.norm(np.maximum(delta, 0.0))
                + min(float(np.max(delta)), 0.0)
            ),
            True,
        )
    if geom_type == int(mujoco.mjtGeom.mjGEOM_ELLIPSOID):
        scaled = float(np.linalg.norm(local / size))
        if scaled <= np.finfo(np.float64).eps:
            return -float(np.min(size)), True
        projected = local / scaled
        distance = float(np.linalg.norm(local - projected))
        return (-distance if scaled < 1.0 else distance), True
    if geom_type == int(mujoco.mjtGeom.mjGEOM_PLANE):
        return float(local[2]), True
    if geom_type == int(mujoco.mjtGeom.mjGEOM_MESH):
        return _mesh_distance(runtime, geom_id, local), True
    radius = max(float(runtime.model.geom_rbound[geom_id]), 1.0e-6)
    return float(np.linalg.norm(local) - radius), False


def _geom_distance(
    runtime: Runtime,
    first: int,
    second: int,
) -> tuple[float, bool]:
    points = np.zeros(6, dtype=np.float64)
    try:
        return (
            float(
                mujoco.mj_geomDistance(
                    runtime.model,
                    runtime.data,
                    first,
                    second,
                    _GEOM_DISTANCE_LIMIT_M,
                    points,
                )
            ),
            True,
        )
    except Exception:
        return -_GEOM_DISTANCE_LIMIT_M, False


def _is_hand_geom(runtime: Runtime, geom_id: int, palm_id: int) -> bool:
    if (
        int(runtime.model.geom_contype[geom_id]) == 0
        and int(runtime.model.geom_conaffinity[geom_id]) == 0
    ):
        return False
    body_id = int(runtime.model.geom_bodyid[geom_id])
    while body_id > 0:
        if body_id == palm_id:
            return True
        body_id = int(runtime.model.body_parentid[body_id])
    return False


def _self_pair_enabled(
    runtime: Runtime,
    first: int,
    second: int,
    excluded: frozenset[int],
) -> bool:
    first_type = int(runtime.model.geom_contype[first])
    first_affinity = int(runtime.model.geom_conaffinity[first])
    second_type = int(runtime.model.geom_contype[second])
    second_affinity = int(runtime.model.geom_conaffinity[second])
    if (first_type & second_affinity) == 0 and (second_type & first_affinity) == 0:
        return False
    first_body = int(runtime.model.geom_bodyid[first])
    second_body = int(runtime.model.geom_bodyid[second])
    if int(runtime.model.body_weldid[first_body]) == int(
        runtime.model.body_weldid[second_body]
    ):
        return False
    if _kinematically_adjacent(runtime.model, first_body, second_body):
        return False
    lower, upper = sorted((first_body, second_body))
    return ((lower << 16) | upper) not in excluded


def _capture_state(runtime: Runtime) -> FloatArray:
    state = np.empty(mujoco.mj_stateSize(runtime.model, _STATE_SPEC), dtype=np.float64)
    mujoco.mj_getState(runtime.model, runtime.data, state, _STATE_SPEC)
    return state


def _restore_state(runtime: Runtime, state: FloatArray) -> None:
    mujoco.mj_setState(runtime.model, runtime.data, state, _STATE_SPEC)
    mujoco.mj_forward(runtime.model, runtime.data)
    mujoco.mj_setState(runtime.model, runtime.data, state, _STATE_SPEC)


def evaluate(runtime: Runtime, scene: Scene, plan: Plan) -> Evaluation:
    geom_ids = tuple(
        int(mujoco.mj_name2id(runtime.model, mujoco.mjtObj.mjOBJ_GEOM, name))
        for name in scene.geom_names
    )
    if not geom_ids or any(index < 0 for index in geom_ids):
        raise ValueError("evaluation object geoms are unavailable")
    body_id = int(
        mujoco.mj_name2id(runtime.model, mujoco.mjtObj.mjOBJ_BODY, scene.body_name)
    )
    if body_id < 0:
        raise ValueError("evaluation object body is unavailable")
    palm_id = int(
        mujoco.mj_name2id(runtime.model, mujoco.mjtObj.mjOBJ_BODY, "palm")
    )
    if palm_id < 0:
        raise ValueError("palm body is unavailable")
    try:
        finger_indices = tuple(FINGER_IDS.index(name) for name in plan.fingers)
    except ValueError as error:
        raise ValueError("plan contains an unknown finger") from error
    active_fingertips = tuple(
        int(
            mujoco.mj_name2id(
                runtime.model,
                mujoco.mjtObj.mjOBJ_GEOM,
                FINGERTIP_GEOMS[index],
            )
        )
        for index in finger_indices
    )
    if not active_fingertips or any(index < 0 for index in active_fingertips):
        raise ValueError("active fingertip geoms are unavailable")
    hand_ids = tuple(
        index
        for index in range(int(runtime.model.ngeom))
        if _is_hand_geom(runtime, index, palm_id)
    )
    excluded = frozenset(
        int(signature) for signature in runtime.model.exclude_signature
    )
    self_pairs = tuple(
        (first, second)
        for position, first in enumerate(hand_ids)
        for second in hand_ids[position + 1 :]
        if _self_pair_enabled(runtime, first, second, excluded)
    )
    state = _capture_state(runtime)
    before_position = np.array(runtime.data.xpos[body_id], copy=True)
    before_rotation = np.array(runtime.data.xmat[body_id], copy=True)
    geometry_supported = True
    contact_errors = np.empty(len(plan.contacts), dtype=np.float64)
    fingertip_clearances = np.empty(len(active_fingertips), dtype=np.float64)
    true_clearance = _GEOM_DISTANCE_LIMIT_M
    path_clearance = _GEOM_DISTANCE_LIMIT_M
    self_clearance = _GEOM_DISTANCE_LIMIT_M
    path_self_clearance = _GEOM_DISTANCE_LIMIT_M
    try:
        for alpha in _adaptive_path_alphas(
            plan.pregrasp,
            plan.joints,
            EXACT_MINIMUM_PATH_SAMPLES,
            EXACT_MAX_PATH_JOINT_STEP_RAD,
        ):
            joints = plan.pregrasp + alpha * (plan.joints - plan.pregrasp)
            runtime.data.qpos[runtime.hand_qpos_indices] = expand_coupled_joints(joints)
            runtime.data.qvel[:] = 0.0
            mujoco.mj_forward(runtime.model, runtime.data)
            object_distances: list[float] = []
            for hand in hand_ids:
                for obj in geom_ids:
                    distance, supported = _geom_distance(runtime, hand, obj)
                    geometry_supported = geometry_supported and supported
                    object_distances.append(distance)
            current_clearance = min(object_distances, default=_GEOM_DISTANCE_LIMIT_M)
            path_clearance = min(path_clearance, current_clearance)
            current_self_values: list[float] = []
            for first, second in self_pairs:
                distance, supported = _geom_distance(runtime, first, second)
                geometry_supported = geometry_supported and supported
                current_self_values.append(distance)
            current_self = np.asarray(current_self_values, dtype=np.float64)
            current_self_clearance = (
                float(np.min(current_self))
                if current_self.size
                else _GEOM_DISTANCE_LIMIT_M
            )
            path_self_clearance = min(path_self_clearance, current_self_clearance)
            if alpha == 1.0:
                true_clearance = current_clearance
                self_clearance = current_self_clearance
        palm_rotation = np.asarray(
            runtime.data.xmat[palm_id], dtype=np.float64
        ).reshape(3, 3)
        palm_position = np.asarray(runtime.data.xpos[palm_id], dtype=np.float64)
        for index, contact in enumerate(plan.contacts):
            point_world = contact.position @ palm_rotation.T + palm_position
            values: list[float] = []
            for geom_id in geom_ids:
                distance, supported = _signed_distance(runtime, geom_id, point_world)
                geometry_supported = geometry_supported and supported
                values.append(abs(distance))
            contact_errors[index] = min(values)
        for index, fingertip in enumerate(active_fingertips):
            values = []
            for geom_id in geom_ids:
                distance, supported = _geom_distance(runtime, fingertip, geom_id)
                geometry_supported = geometry_supported and supported
                values.append(distance)
            fingertip_clearances[index] = min(values)
    finally:
        _restore_state(runtime, state)
    unchanged = bool(
        np.allclose(runtime.data.xpos[body_id], before_position, atol=1.0e-12)
        and np.allclose(runtime.data.xmat[body_id], before_rotation, atol=1.0e-12)
    )
    return Evaluation(
        contact_errors,
        fingertip_clearances,
        float(true_clearance),
        unchanged,
        float(path_clearance),
        float(self_clearance),
        float(path_self_clearance),
        geometry_supported,
    )


__all__ = [
    "EXACT_CLEARANCE_TOLERANCE_M",
    "EXACT_CONTACT_TOLERANCE_M",
    "EXACT_FINGERTIP_GAP_M",
    "EXACT_MAX_PATH_JOINT_STEP_RAD",
    "EXACT_MINIMUM_PATH_SAMPLES",
    "EXACT_PATH_SAMPLES",
    "Evaluation",
    "SELF_COLLISION_TOLERANCE_M",
    "evaluate",
]