File size: 5,631 Bytes
7399b6f
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
# Copyright (c) 2022-2025, The Isaac Lab Project Developers.
# All rights reserved.
#
# SPDX-License-Identifier: BSD-3-Clause

"""A cuRobo-free PRM (Probabilistic Roadmap) motion planner in end-effector task space.

This is the motion-planning primitive for the Galbot dual-parallel TAMP integration
when cuRobo cannot be compiled. It plans a collision-free polyline for the end
effector between waypoints, checking straight-line edges against axis-aligned box
obstacles (oven body, table, fridge, ...) inflated by an EEF clearance radius.

The env's RMPFlow controller turns each EEF waypoint into joint commands, so a
task-space roadmap over reachable EEF positions is sufficient to drive the arm
around obstacles; grasp orientations come from :mod:`mesh_grasp`.

Pure numpy + scipy (KD-tree). Deterministic given ``seed``.
"""

from __future__ import annotations

from dataclasses import dataclass, field

import numpy as np


@dataclass
class Box:
    """Axis-aligned box obstacle in world/env-local frame."""

    center: np.ndarray
    half: np.ndarray  # half extents (x,y,z)

    def contains(self, p: np.ndarray, pad: float = 0.0) -> bool:
        d = np.abs(np.asarray(p) - self.center)
        return bool(np.all(d <= (self.half + pad)))

    def segment_hits(self, a: np.ndarray, b: np.ndarray, pad: float = 0.0, steps: int = 24) -> bool:
        """Conservative segment-box test by dense sampling (robust, simple)."""
        for t in np.linspace(0.0, 1.0, steps):
            if self.contains(a + t * (b - a), pad):
                return True
        return False


@dataclass
class PRM:
    bounds_lo: np.ndarray
    bounds_hi: np.ndarray
    obstacles: list[Box] = field(default_factory=list)
    clearance: float = 0.05      # EEF sphere radius for collision inflation
    num_samples: int = 300
    k: int = 10
    seed: int = 0

    def _free_point(self, p: np.ndarray) -> bool:
        if np.any(p < self.bounds_lo) or np.any(p > self.bounds_hi):
            return False
        return not any(o.contains(p, self.clearance) for o in self.obstacles)

    def _free_edge(self, a: np.ndarray, b: np.ndarray) -> bool:
        return not any(o.segment_hits(a, b, self.clearance) for o in self.obstacles)

    def plan(self, start: np.ndarray, goal: np.ndarray) -> np.ndarray | None:
        """Return an (M,3) collision-free polyline from start to goal, or None."""
        from scipy.spatial import cKDTree

        start = np.asarray(start, dtype=np.float64)
        goal = np.asarray(goal, dtype=np.float64)
        rng = np.random.default_rng(self.seed)

        # direct connection first
        if self._free_point(start) and self._free_point(goal) and self._free_edge(start, goal):
            return np.stack([start, goal]).astype(np.float32)

        # sample free nodes
        nodes = [start, goal]
        tries = 0
        while len(nodes) < self.num_samples + 2 and tries < self.num_samples * 30:
            tries += 1
            p = rng.uniform(self.bounds_lo, self.bounds_hi)
            if self._free_point(p):
                nodes.append(p)
        nodes = np.array(nodes)
        n = len(nodes)
        if n < 3:
            return None

        # k-NN graph with collision-free edges
        tree = cKDTree(nodes)
        adj: dict[int, list[tuple[int, float]]] = {i: [] for i in range(n)}
        for i in range(n):
            dists, idxs = tree.query(nodes[i], k=min(self.k + 1, n))
            for d, j in zip(np.atleast_1d(dists), np.atleast_1d(idxs)):
                j = int(j)
                if j == i:
                    continue
                if self._free_edge(nodes[i], nodes[j]):
                    adj[i].append((j, float(d)))

        # Dijkstra start(0) -> goal(1)
        import heapq

        dist = {0: 0.0}
        prev: dict[int, int] = {}
        pq = [(0.0, 0)]
        while pq:
            d, u = heapq.heappop(pq)
            if u == 1:
                break
            if d > dist.get(u, np.inf):
                continue
            for v, w in adj[u]:
                nd = d + w
                if nd < dist.get(v, np.inf):
                    dist[v] = nd
                    prev[v] = u
                    heapq.heappush(pq, (nd, v))
        if 1 not in dist:
            return None
        # backtrace
        path = [1]
        while path[-1] != 0:
            path.append(prev[path[-1]])
        path = path[::-1]
        return nodes[path].astype(np.float32)


def shortcut(path: np.ndarray, obstacles: list[Box], clearance: float, iters: int = 100,
             seed: int = 0) -> np.ndarray:
    """Randomized path shortcutting: greedily replace subpaths with straight edges."""
    if path is None or len(path) < 3:
        return path
    rng = np.random.default_rng(seed)
    pts = [p for p in path]
    for _ in range(iters):
        if len(pts) < 3:
            break
        i = rng.integers(0, len(pts) - 2)
        j = rng.integers(i + 2, len(pts))
        a, b = pts[i], pts[j]
        if not any(o.segment_hits(a, b, clearance) for o in obstacles):
            pts = pts[: i + 1] + pts[j:]
    return np.array(pts, dtype=np.float32)


def resample_polyline(path: np.ndarray, n: int) -> np.ndarray:
    """Arc-length resample an (M,3) polyline to exactly n points."""
    path = np.asarray(path, dtype=np.float64)
    seg = np.linalg.norm(np.diff(path, axis=0), axis=1)
    cum = np.concatenate([[0.0], np.cumsum(seg)])
    total = cum[-1] if cum[-1] > 1e-9 else 1.0
    ts = np.linspace(0.0, total, n)
    out = np.empty((n, 3))
    for d in range(3):
        out[:, d] = np.interp(ts, cum, path[:, d])
    return out.astype(np.float32)