diff --git a/skrobot/planner/sampling.py b/skrobot/planner/sampling.py new file mode 100644 index 00000000..74ad9608 --- /dev/null +++ b/skrobot/planner/sampling.py @@ -0,0 +1,343 @@ +"""Sampling-based planning in joint space. + +An optimizer such as :func:`skrobot.planner.collision_aware.plan_trajectory` +improves the trajectory it is started from and stays in its basin: seeded +with a straight joint-space path that runs through an obstacle, it can end +up pressed against that obstacle rather than around it. A sampling planner +has the opposite profile -- it finds a way around, in whatever homotopy +class, but the way is jagged -- so the two are run in sequence: this module +finds a collision-free path, the optimizer smooths it. + +:func:`rrt_connect` is the bidirectional RRT of Kuffner and LaValle [1]_, +which grows one tree from the start and one from the goal and tries to +join them after every extension; :func:`shortcut` removes the detours a +random tree leaves behind. + +References +---------- +.. [1] J. J. Kuffner and S. M. LaValle. "RRT-Connect: An Efficient + Approach to Single-Query Path Planning." IEEE International Conference + on Robotics and Automation (ICRA), 2000. +""" + +import time + +import numpy as np + + +_BATCH_CHUNK = 8 + + +def _segment_is_valid(is_valid, a, b, resolution, is_valid_batch=None): + """Whether the straight motion from ``a`` to ``b`` is valid throughout. + + Checked at points no more than ``resolution`` apart in max-norm, end + excluded (the caller has checked the endpoints). With + ``is_valid_batch`` every point is checked in one call. + """ + n = int(np.ceil(np.max(np.abs(b - a)) / resolution)) + if n <= 1: + return True + if is_valid_batch is not None: + # A few points per call: most candidate motions a shortcut tries + # fail, and a long one checked all at once would pay for every + # point where the first bad one settles it. The coarse points + # (the midpoint, the quarters) go first, so a motion through an + # obstacle usually fails on the first call. + fractions = np.arange(1, n) / float(n) + order = np.argsort(np.abs(fractions - 0.5) + 1e-3 * fractions, + kind='stable') + fractions = fractions[order] + for chunk in range(0, len(fractions), _BATCH_CHUNK): + part = fractions[chunk:chunk + _BATCH_CHUNK] + points = a[None, :] + (b - a)[None, :] * part[:, None] + if not np.all(is_valid_batch(points)): + return False + return True + for i in range(1, n): + if not is_valid(a + (b - a) * (i / float(n))): + return False + return True + + +class _Tree(object): + + def __init__(self, root): + self.nodes = [np.asarray(root, dtype=np.float64)] + self.parents = [-1] + + def nearest(self, q): + nodes = np.asarray(self.nodes) + return int(np.argmin(np.max(np.abs(nodes - q), axis=1))) + + def add(self, q, parent): + self.nodes.append(np.asarray(q, dtype=np.float64)) + self.parents.append(int(parent)) + return len(self.nodes) - 1 + + def path_to_root(self, index): + path = [] + while index >= 0: + path.append(self.nodes[index]) + index = self.parents[index] + return path + + +def rrt_connect(start, goal, is_valid, lower, upper, step=0.2, + resolution=0.05, max_iterations=5000, time_limit=5.0, + goal_bias=0.1, line_bias=0.5, line_sigma=0.5, seed=0, + is_valid_batch=None): + """Find a collision-free joint-space path with bidirectional RRT. + + Parameters + ---------- + start : array-like + Joint configuration, valid. + goal : array-like + A valid joint configuration, or several as an ``(n, n_joints)`` + array -- the goal tree is then rooted at all of them and the + path ends at whichever the start tree reaches, which on an arm + with several ways to reach a pose finds a way far more often. + is_valid : callable + ``q -> bool``, whether a configuration is collision-free. + lower, upper : array-like + Joint limits the samples are drawn from. + step : float + How far (max-norm, rad) a tree grows toward a sample at a time. + resolution : float + Spacing (max-norm) of the validity checks along an extension. + This is what the returned path is checked at and all it can be + trusted to: a violation narrower than ``resolution`` can pass + between two check points, so a caller that needs more than that + must either check the path itself at a finer spacing or give + ``is_valid`` a margin of its own. + max_iterations : int + Sample budget. + time_limit : float + Wall-clock budget in seconds; the search stops at whichever of + the two budgets runs out first. + goal_bias : float + Fraction of samples that aim straight at the other tree's root. + line_bias : float + Fraction of the remaining samples drawn around the straight + line from start to goal -- a point on it plus Gaussian noise of + ``line_sigma`` per joint -- rather than uniformly over the joint + limits. An arm's detour around an obstacle lies near that line; + most of joint space does not, and uniform samples there are + wasted. + line_sigma : float + Spread (rad, or m) of the samples around the line. + seed : int + Random seed, so a run is reproducible. + + Returns + ------- + list[numpy.ndarray] or None + Configurations from ``start`` to the goal reached whose + consecutive straight motions are valid, or ``None`` when no + path was found within the budgets. + """ + start = np.asarray(start, dtype=np.float64) + goals = np.atleast_2d(np.asarray(goal, dtype=np.float64)) + lower = np.asarray(lower, dtype=np.float64) + upper = np.asarray(upper, dtype=np.float64) + rng = np.random.RandomState(seed) + goal_tree = _Tree(goals[0]) + for extra in goals[1:]: + goal_tree.add(extra, -1) + trees = (_Tree(start), goal_tree) + deadline = time.time() + float(time_limit) + + def extend(tree, q_target): + """Grow ``tree`` toward ``q_target``; return (status, node index). + + ``status`` is 'reached' when the target itself was added, + 'advanced' when a step toward it was, 'trapped' otherwise. + """ + near = tree.nearest(q_target) + q_near = tree.nodes[near] + delta = q_target - q_near + span = float(np.max(np.abs(delta))) + if span <= step: + q_new, status = q_target, 'reached' + else: + q_new, status = q_near + delta * (step / span), 'advanced' + if is_valid_batch is not None: + n_sub = int(np.ceil(span / resolution)) + fractions = np.arange(1, max(n_sub, 1) + 1) / float(max(n_sub, 1)) + points = q_near[None, :] + (q_new - q_near)[None, :] \ + * fractions[:, None] + if not np.all(is_valid_batch(points)): + return 'trapped', None + elif not is_valid(q_new) or not _segment_is_valid( + is_valid, q_near, q_new, resolution): + return 'trapped', None + return status, tree.add(q_new, near) + + def connect(tree, q_target): + status = 'advanced' + index = None + while status == 'advanced': + status, index = extend(tree, q_target) + return status, index + + for iteration in range(int(max_iterations)): + if time.time() > deadline: + return None + a, b = trees[iteration % 2], trees[(iteration + 1) % 2] + if rng.uniform() < goal_bias: + roots = [i for i, p in enumerate(b.parents) if p < 0] + sample = b.nodes[roots[rng.randint(len(roots))]] + elif rng.uniform() < line_bias: + end = goals[rng.randint(len(goals))] + along = start + (end - start) * rng.uniform() + sample = np.clip(along + rng.normal(0.0, line_sigma, start.shape), + lower, upper) + else: + sample = rng.uniform(lower, upper) + status, index_a = extend(a, sample) + if status == 'trapped': + continue + status, index_b = connect(b, a.nodes[index_a]) + if status != 'reached': + continue + path_a = a.path_to_root(index_a)[::-1] + path_b = b.path_to_root(index_b) + path = path_a + path_b[1:] + if a is trees[1]: + path = path[::-1] + return path + return None + + +def shortcut(path, is_valid, resolution=0.05, iterations=100, seed=0, + is_valid_batch=None): + """Remove detours from a path by straight-line shortcuts. + + Repeatedly picks two points on the path and replaces what lies + between them by the straight motion when that motion is valid. + + Parameters + ---------- + path : list[numpy.ndarray] + Configurations whose consecutive straight motions are valid. + is_valid : callable + ``q -> bool``. + resolution : float + Spacing (max-norm) of the validity checks along a shortcut. + iterations : int + Shortcut attempts. + seed : int + Random seed. + + Returns + ------- + list[numpy.ndarray] + The shortened path, same endpoints. + """ + path = [np.asarray(q, dtype=np.float64) for q in path] + rng = np.random.RandomState(seed) + for _ in range(int(iterations)): + if len(path) < 3: + break + i, j = sorted(rng.choice(len(path), 2, replace=False)) + if j - i < 2: + continue + if _segment_is_valid(is_valid, path[i], path[j], resolution, + is_valid_batch=is_valid_batch): + path = path[:i + 1] + path[j:] + return path + + +def resample(path, n_points): + """Spread ``n_points`` configurations evenly along a path. + + Even in joint-space arc length (max-norm), so the seed a planner + hands to an optimizer has comparable steps everywhere. + + Parameters + ---------- + path : list[numpy.ndarray] + Configurations; consecutive ones are joined by straight motions. + n_points : int + Number of configurations wanted, endpoints included. + + Returns + ------- + numpy.ndarray + ``(n_points, n_joints)``, first and last rows the path's ends. + """ + path = np.asarray(path, dtype=np.float64) + if len(path) == 1: + return np.tile(path, (int(n_points), 1)) + lengths = np.max(np.abs(np.diff(path, axis=0)), axis=1) + cumulative = np.concatenate([[0.0], np.cumsum(lengths)]) + total = float(cumulative[-1]) + targets = np.linspace(0.0, total, int(n_points)) + out = np.empty((int(n_points), path.shape[1])) + for k, s in enumerate(targets): + i = int(min(np.searchsorted(cumulative, s, side='right') - 1, + len(lengths) - 1)) + seg = lengths[i] + t = 0.0 if seg <= 0.0 else (s - cumulative[i]) / seg + out[k] = path[i] + (path[i + 1] - path[i]) * min(max(t, 0.0), 1.0) + out[0] = path[0] + out[-1] = path[-1] + return out + + +def smooth_path(path, is_valid, n_points, rounds=3, resolution=0.05, + iterations=200, seed=0, is_valid_batch=None): + """Shortcut a path, resample it and round its corners, all checked. + + The sampling planner's own finishing pass, with no optimizer: random + shortcuts remove the tree's detours, the path is resampled to + ``n_points`` evenly spaced waypoints, and each interior waypoint is + then pulled toward the mean of its neighbours whenever the moved + waypoint and both motions around it stay valid. Every step keeps + the path valid under ``is_valid``; nothing is left to a solver. + + Parameters + ---------- + path : numpy.ndarray + ``(m, n_joints)`` waypoints from start to goal. + is_valid : callable + ``f(q) -> bool``. + n_points : int + Waypoints in the result. + rounds : int + Passes of the corner rounding. + resolution : float + Largest max-norm step between validity checks along a motion. + iterations : int + Shortcut attempts. + seed : int + Seed of the shortcut's random pairs. + is_valid_batch : callable or None + ``f(points) -> (m,) bools`` checking many configurations in one + call; used for every motion when given. + + Returns + ------- + numpy.ndarray + ``(n_points, n_joints)`` waypoints, endpoints unchanged. + """ + path = shortcut(np.asarray(path, dtype=np.float64), is_valid, + resolution=resolution, iterations=iterations, seed=seed, + is_valid_batch=is_valid_batch) + path = resample(path, n_points) + for _ in range(int(rounds)): + rounded = path.copy() + for i in range(1, len(path) - 1): + candidate = 0.25 * path[i - 1] + 0.5 * path[i] + 0.25 * path[i + 1] + if _segment_is_valid(is_valid, rounded[i - 1], candidate, + resolution, is_valid_batch=is_valid_batch) \ + and _segment_is_valid(is_valid, candidate, path[i + 1], + resolution, + is_valid_batch=is_valid_batch) \ + and (is_valid_batch(candidate[None])[0] + if is_valid_batch is not None + else is_valid(candidate)): + rounded[i] = candidate + path = rounded + return path diff --git a/tests/skrobot_tests/planner_tests/test_sampling.py b/tests/skrobot_tests/planner_tests/test_sampling.py new file mode 100644 index 00000000..97ab67fa --- /dev/null +++ b/tests/skrobot_tests/planner_tests/test_sampling.py @@ -0,0 +1,82 @@ +import unittest + +import numpy as np +from numpy import testing + +from skrobot.planner.sampling import resample +from skrobot.planner.sampling import rrt_connect +from skrobot.planner.sampling import shortcut + + +def _wall_with_gap(q): + """A 2-D world: a wall at x in [0.4, 0.6] with a gap at y > 1.5.""" + x, y = q + return not (0.4 <= x <= 0.6 and y <= 1.5) + + +class TestRRTConnect(unittest.TestCase): + + def test_finds_a_way_through_the_gap(self): + start = np.array([0.0, 0.0]) + goal = np.array([1.0, 0.0]) + lower = np.array([-0.5, -0.5]) + upper = np.array([1.5, 2.0]) + path = rrt_connect(start, goal, _wall_with_gap, lower, upper, + step=0.2, resolution=0.05, seed=1) + self.assertIsNotNone(path) + testing.assert_allclose(path[0], start) + testing.assert_allclose(path[-1], goal) + # Every configuration, and every step between two, is valid at + # the resolution the search checked at. Sampling finer than that + # is asking for a guarantee the algorithm does not make: a + # violation narrower than one check spacing can pass between two + # check points, and this world has a hard boundary to graze. + for a, b in zip(path[:-1], path[1:]): + n = max(int(np.ceil(np.max(np.abs(b - a)) / 0.05)), 1) + for i in range(n + 1): + self.assertTrue(_wall_with_gap(a + (b - a) * (i / float(n)))) + # It went through the gap, not the wall. + self.assertGreater(max(q[1] for q in path), 1.5) + + # Shortcutting keeps the endpoints and validity and never lengthens. + length = sum(np.max(np.abs(b - a)) for a, b in zip(path[:-1], path[1:])) + short = shortcut(path, _wall_with_gap, resolution=0.05, seed=1) + testing.assert_allclose(short[0], start) + testing.assert_allclose(short[-1], goal) + short_length = sum(np.max(np.abs(b - a)) + for a, b in zip(short[:-1], short[1:])) + self.assertLessEqual(short_length, length + 1e-9) + for a, b in zip(short[:-1], short[1:]): + n = max(int(np.ceil(np.max(np.abs(b - a)) / 0.05)), 1) + for i in range(n + 1): + self.assertTrue(_wall_with_gap(a + (b - a) * (i / float(n)))) + + def test_gives_up_within_the_budget(self): + # No gap: nothing can get across. + def wall(q): + return not (0.4 <= q[0] <= 0.6) + path = rrt_connect(np.array([0.0, 0.0]), np.array([1.0, 0.0]), wall, + np.array([-0.5, -0.5]), np.array([1.5, 2.0]), + max_iterations=300, time_limit=2.0, seed=0) + self.assertIsNone(path) + + +class TestResample(unittest.TestCase): + + def test_spreads_points_evenly_and_keeps_the_ends(self): + path = [np.array([0.0, 0.0]), np.array([1.0, 0.0]), + np.array([1.0, 3.0])] + out = resample(path, 5) + self.assertEqual(out.shape, (5, 2)) + testing.assert_allclose(out[0], path[0]) + testing.assert_allclose(out[-1], path[-1]) + # Arc length (max-norm) is 4; points every 1.0 along it. + testing.assert_allclose(out[1], [1.0, 0.0]) + testing.assert_allclose(out[2], [1.0, 1.0]) + testing.assert_allclose(out[3], [1.0, 2.0]) + steps = np.max(np.abs(np.diff(out, axis=0)), axis=1) + testing.assert_allclose(steps, 1.0) + + +if __name__ == '__main__': + unittest.main()