src/main.py - master - koda/FiVI

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

import argparse
import json
import logging
from pathlib import Path
import sys
from time import perf_counter_ns, sleep
from typing import NoReturn, Sequence

import numpy as np

from .evaluation import Evaluation, evaluate
from .ik import JointLimits, expand_coupled_joints
from .perception import Observation, Observer
from .planner import Plan, Planner, Result
from .runtime import Runtime
from .scene import Scene, Variation, build_scene, resolve_asset
from .viewer import add


class _ArgumentParser(argparse.ArgumentParser):
    def error(self, message: str) -> NoReturn:
        raise ValueError(message)


def _parser() -> _ArgumentParser:
    parser = _ArgumentParser(description="RGB-D static grasp pose planner")
    parser.add_argument("--object", default="sphere", help="cube, sphere, cylinder, irregular, or object XML")
    parser.add_argument("--offset", nargs=3, type=float, metavar=("X", "Y", "Z"), default=(0.0, 0.0, 0.0))
    parser.add_argument("--roll", type=float, default=0.0)
    parser.add_argument("--pitch", type=float, default=0.0)
    parser.add_argument("--yaw", type=float, default=0.0)
    parser.add_argument("--scale", type=float, default=1.0)
    parser.add_argument("--viewer", action="store_true")
    parser.add_argument("--debug-view", action="store_true")
    parser.add_argument("--json", action="store_true")
    parser.add_argument("--no-evaluation", action="store_true")
    return parser


def _limits(runtime: Runtime) -> JointLimits:
    ranges = np.asarray(
        runtime.model.actuator_ctrlrange[runtime.hand_actuator_ids],
        dtype=np.float64,
    )
    return JointLimits(ranges[:, 0], ranges[:, 1])


def _validate(
    runtime: Runtime,
    scene: Scene,
    result: Result,
) -> tuple[Result, Evaluation | None, int, int]:
    measured: list[tuple[Plan, Evaluation]] = []
    for plan in result.validation_candidates:
        try:
            measured.append((plan, evaluate(runtime, scene, plan)))
        except (RuntimeError, ValueError) as error:
            logging.warning("exact evaluation rejected a plan: %s", error)
    accepted = tuple(item for item in measured if item[1].accepted)
    validated = result.accept_validated(item[0] for item in accepted)
    if accepted:
        return (
            validated,
            accepted[0][1],
            len(result.validation_candidates),
            len(accepted),
        )
    return validated, None, len(result.validation_candidates), 0


def _plan_payload(plan: Plan) -> dict[str, object]:
    joints = expand_coupled_joints(plan.joints)
    pregrasp = expand_coupled_joints(plan.pregrasp)
    return {
        "topology": plan.topology.value,
        "fingers": plan.fingers,
        "contacts_hand_m": tuple(
            tuple(float(value) for value in item.position) for item in plan.contacts
        ),
        "fingertips_hand_m": tuple(
            tuple(float(value) for value in item) for item in plan.fingertips
        ),
        "joint_positions_reduced_rad": tuple(float(value) for value in plan.joints),
        "joint_positions_full_rad": tuple(float(value) for value in joints),
        "pregrasp_reduced_rad": tuple(float(value) for value in plan.pregrasp),
        "pregrasp_full_rad": tuple(float(value) for value in pregrasp),
        "surface_errors_mm": tuple(
            float(value * 1000.0) for value in plan.surface_errors_m
        ),
        "fingertip_errors_mm": tuple(
            float(value * 1000.0) for value in plan.errors_m
        ),
        "estimated_clearance_mm": 1000.0 * plan.estimated_clearance_m,
        "self_clearance_mm": 1000.0 * plan.self_clearance_m,
        "sampled_path_clearance_mm_estimated": (
            1000.0 * plan.sampled_path_clearance_m
        ),
        "sampled_path_self_clearance_mm_estimated": (
            1000.0 * plan.sampled_path_self_clearance_m
        ),
        "score": plan.score,
        "ik_time_ms": plan.ik_time_ms,
        "collision_time_ms": plan.collision_time_ms,
    }


def _evaluation_payload(measured: Evaluation) -> dict[str, object]:
    return {
        "rendered_surface_errors_mm": tuple(
            float(value * 1000.0) for value in measured.contact_errors_m
        ),
        "fingertip_clearances_mm": tuple(
            float(value * 1000.0) for value in measured.fingertip_clearances_m
        ),
        "true_clearance_mm": 1000.0 * measured.true_clearance_m,
        "path_clearance_mm": 1000.0 * measured.path_clearance_m,
        "self_clearance_mm_exact": 1000.0 * measured.self_clearance_m,
        "path_self_clearance_mm": 1000.0 * measured.path_self_clearance_m,
        "sampled_path_clearance_mm": (
            1000.0 * measured.sampled_path_clearance_m
        ),
        "sampled_path_self_clearance_mm": (
            1000.0 * measured.sampled_path_self_clearance_m
        ),
        "object_pose_unchanged": measured.object_pose_unchanged,
        "geometry_supported": measured.geometry_supported,
    }


def run(
    object_name: str,
    variation: Variation = Variation(),
    *,
    exact_evaluation: bool = True,
) -> tuple[Runtime, tuple[Scene, Observation], Result, dict[str, object]]:
    """Assemble and execute the static pipeline; the returned runtime must be closed."""

    root = Path(__file__).resolve().parents[1]
    object_xml = resolve_asset(object_name, root / "models" / "objects")
    scene = build_scene(root / "models" / "hand" / "hand.xml", object_xml, variation)
    runtime = Runtime(scene.runtime)
    try:
        observation = Observer(runtime, scene).observe()
        limits = _limits(runtime)
        neutral = np.clip(np.zeros(16, dtype=np.float64), limits.lower, limits.upper)
        planner = Planner(runtime, observation, limits, neutral)
        result = planner.search()
        approximate_feasible = result.feasible
        measured: Evaluation | None = None
        evaluated_plans = 0
        accepted_plans = 0
        validation_time_ms = 0.0
        if exact_evaluation:
            validation_started = perf_counter_ns()
            result, measured, evaluated_plans, accepted_plans = _validate(
                runtime,
                scene,
                result,
            )
            validation_time_ms = (perf_counter_ns() - validation_started) / 1.0e6
        validation_status = (
            "passed"
            if exact_evaluation and result.plans
            else "failed"
            if exact_evaluation and approximate_feasible > 0
            else "not_applicable"
            if exact_evaluation
            else "skipped"
        )
        payload: dict[str, object] = {
            "schema_version": 1,
            "object": object_name,
            "variation": {
                "offset_m": tuple(float(value) for value in variation.offset_m),
                "rotation_deg": tuple(
                    float(value) for value in variation.rotation_deg
                ),
                "scale": float(variation.scale),
            },
            "surface": result.shape.value,
            "generated": result.generated,
            "feasible": result.feasible,
            "approximate_feasible": approximate_feasible,
            "failure": None if result.failure is None else result.failure.value,
            "planning_time_ms": result.planning_time_ms,
            "perception": {
                "timestamp_s": observation.timestamp_s,
                "camera_name": observation.camera_name,
                "position_hand_m": tuple(
                    float(value) for value in observation.position_hand
                ),
                "orientation_hand": tuple(
                    tuple(float(value) for value in row)
                    for row in observation.orientation_hand
                ),
                "dimensions_m": tuple(
                    float(value) for value in observation.dimensions_m
                ),
                "confidence": observation.confidence,
                "fit_residual_mm": 1000.0 * observation.fit_residual_m,
                "uncertainty_margin_mm": 1000.0 * observation.uncertainty_margin_m,
                "visible_surface_fraction": observation.visible_surface_fraction,
            },
            "validation": {
                "requested": exact_evaluation,
                "status": validation_status,
                "evaluated_plans": evaluated_plans,
                "accepted_plans": accepted_plans,
                "time_ms": validation_time_ms,
            },
        }
        if result.plans:
            plan = result.plans[0]
            runtime.set_joint_positions(expand_coupled_joints(plan.joints))
            payload.update(_plan_payload(plan))
            if measured is not None:
                payload.update(_evaluation_payload(measured))
        return runtime, (scene, observation), result, payload
    except BaseException:
        runtime.close()
        raise


def main(argv: Sequence[str] | None = None) -> int:
    """Run the single supported CLI without import-time side effects."""

    source_arguments = tuple(sys.argv[1:] if argv is None else argv)
    json_requested = "--json" in source_arguments
    logging.basicConfig(level=logging.INFO, format="%(levelname)s %(message)s")
    runtime: Runtime | None = None
    try:
        arguments = _parser().parse_args(source_arguments)
        if bool(arguments.json) and (
            bool(arguments.viewer) or bool(arguments.debug_view)
        ):
            raise ValueError("--json cannot be combined with interactive viewing")
        variation = Variation(
            (
                float(arguments.offset[0]),
                float(arguments.offset[1]),
                float(arguments.offset[2]),
            ),
            (float(arguments.roll), float(arguments.pitch), float(arguments.yaw)),
            float(arguments.scale),
        )
        runtime, context, result, payload = run(
            str(arguments.object),
            variation,
            exact_evaluation=not bool(arguments.no_evaluation),
        )
        if bool(arguments.json):
            print(json.dumps(payload, allow_nan=False, indent=2, sort_keys=True))
        else:
            for key, value in payload.items():
                print(f"{key}: {value}")
        if bool(arguments.viewer) or bool(arguments.debug_view):
            if not result.plans:
                return 2
            import mujoco.viewer as mj_viewer  # type: ignore[import-untyped]

            scene, observation = context
            del scene
            with mj_viewer.launch_passive(runtime.model, runtime.data) as handle:
                add(handle, runtime, observation, result.plans[0], debug=bool(arguments.debug_view))
                while handle.is_running():
                    handle.sync()
                    sleep(0.02)
        return 0 if result.plans else 2
    except (OSError, RuntimeError, ValueError) as error:
        if json_requested:
            print(
                json.dumps(
                    {
                        "schema_version": 1,
                        "failure": "PIPELINE_ERROR",
                        "error": str(error),
                    },
                    allow_nan=False,
                    indent=2,
                    sort_keys=True,
                )
            )
        logging.error("%s", error)
        return 2
    finally:
        if runtime is not None:
            runtime.close()


if __name__ == "__main__":
    raise SystemExit(main())