"""OMPL and MuJoCo joint-space motion-planning utilities.
This module connects MuJoCo collision checking with OMPL geometric
planners. It provides helpers for creating OMPL state spaces, converting
between OMPL states and NumPy vectors, configuring planners, validating
paths, repairing invalid path segments, and planning collision-free robot
motions directly from RobotBlockSet MuJoCo robot objects.
Existing collision-free paths can also be inserted as incumbent OMPL
solutions and refined further.
Notes
-----
MuJoCo ``MjData`` is not thread-safe. :class:`JointPathPlanner` keeps a
private copy for state-validity checks and refreshes it when the scene
model/data are updated.
Some OMPL planner classes are optional in Python builds. Unsupported or
unavailable planner names raise ``ValueError`` with a descriptive message.
Copyright (c) 2026 Jozef Stefan Institute
Authors: Leon Zlajpah.
"""
from __future__ import annotations
# pyright: reportMissingImports=false
import numpy as np
from copy import copy
from typing import Callable, Iterable, Optional, Sequence, Tuple, TYPE_CHECKING, Union
try:
import mujoco
except Exception as e:
raise ModuleNotFoundError(f"{e}\nMuJoCo not installed. \nYou can install MuJoCo through pip:\n pip install mujoco") from None
try:
from ompl import base as ob, geometric as og
except Exception as e:
raise ModuleNotFoundError(f"{e}\nOpen Motion Planning Library (OMPL) not installed. \nYou can install OMPL through pip:\n pip install ompl") from None
from robotblockset.mujoco.tools_pymujoco import check_object_weld, get_robot_joints_data, get_body_descendants, get_geoms_of_body, geoms_of_welded_objects, weld_all_objects
from robotblockset.tools import check_option
if TYPE_CHECKING:
from robotblockset.robots import robot
StateValidityFn = Callable[[ob.State], bool]
VectorLike = Union[np.ndarray, Sequence[float]]
MatrixLike = Union[np.ndarray, Sequence[Sequence[float]]]
PathLike = Union[np.ndarray, Sequence[float], Sequence[Sequence[float]]]
PathResult = Tuple[Optional[np.ndarray], Optional[np.ndarray]]
CollisionMarginState = Optional[Tuple[Optional[np.ndarray], Union[np.ndarray, dict[int, float]]]]
ContactGeomPairLike = Iterable[int]
[docs]
def make_state_space(bounds: MatrixLike) -> ob.RealVectorStateSpace:
"""Create an OMPL real-vector state space from numeric bounds.
Parameters
----------
bounds : array_like, shape (n, 2)
Lower and upper limits for each state-space dimension.
Returns
-------
ompl.base.RealVectorStateSpace
State space with the supplied bounds.
Raises
------
ValueError
If ``bounds`` is not two-dimensional with two columns, has no
dimensions, or contains a lower bound greater than its upper bound.
"""
arr = np.asarray(bounds, dtype=float)
if arr.ndim != 2 or arr.shape[1] != 2:
raise ValueError("bounds must have shape (n, 2)")
if arr.shape[0] == 0:
raise ValueError("bounds must define at least one dimension")
if np.any(arr[:, 0] > arr[:, 1]):
i = int(np.argmax(arr[:, 0] > arr[:, 1]))
raise ValueError(f"low > high for dimension {i}: {arr[i]}")
dim = int(arr.shape[0])
space = ob.RealVectorStateSpace(dim)
rb = ob.RealVectorBounds(dim)
for i, (lo, hi) in enumerate(arr):
rb.setLow(i, float(lo))
rb.setHigh(i, float(hi))
space.setBounds(rb)
return space
[docs]
def vec_to_state(space: ob.StateSpace, v: VectorLike) -> ob.State:
"""Convert a vector of coordinates to an OMPL state.
Parameters
----------
space : ompl.base.StateSpace
Target OMPL state space.
v : sequence of float
State coordinates. The length must equal ``space.getDimension()``.
Returns
-------
ompl.base.State
Newly allocated state populated with ``v``.
Raises
------
ValueError
If ``len(v)`` does not match the space dimension.
"""
dim = space.getDimension()
if len(v) != dim:
raise ValueError(f"Expected vector of length {dim}, got {len(v)}")
try:
s = ob.State(space)
except TypeError:
s = space.allocState()
for i, val in enumerate(v):
s[i] = float(val)
return s
[docs]
def state_to_vec(s: ob.State, n: int) -> np.ndarray:
"""Convert an OMPL state to a NumPy vector.
Parameters
----------
s : ompl.base.State
State to convert.
n : int
Number of scalar coordinates to read from ``s``.
Returns
-------
ndarray, shape (n,)
State coordinates.
"""
return np.fromiter((s[i] for i in range(n)), dtype=float, count=n)
def _as_ompl_state(state: ob.State) -> ob.State:
"""Return a state object accepted by OMPL methods.
Parameters
----------
state : ompl.base.State
OMPL state or callable state wrapper produced by some OMPL Python
bindings.
Returns
-------
ompl.base.State
State object in the form expected by OMPL path and validity methods.
"""
return state() if callable(state) else state
def _append_state(path: og.PathGeometric, state: ob.State) -> None:
"""Append a state to an OMPL geometric path.
Parameters
----------
path : ompl.geometric.PathGeometric
Path to modify.
state : ompl.base.State
State to append. Callable wrappers are unwrapped for compatibility with
multiple OMPL Python bindings.
Returns
-------
None
"""
path.append(_as_ompl_state(state))
def _set_state_validity_checker(ss: og.SimpleSetup, validity_fn: StateValidityFn) -> None:
"""Set a Python state-validity callback on an OMPL ``SimpleSetup``.
Parameters
----------
ss : ompl.geometric.SimpleSetup
Planning setup to configure.
validity_fn : callable
Function receiving an OMPL state and returning ``True`` for valid
states.
Returns
-------
None
"""
checker_factory = getattr(ob, "StateValidityCheckerFn", None)
checker = checker_factory(validity_fn) if checker_factory is not None else validity_fn
ss.setStateValidityChecker(checker)
def _path_geometric_from_waypoints(space: ob.StateSpace, si: ob.SpaceInformation, waypoints: MatrixLike) -> og.PathGeometric:
"""Build an OMPL geometric path from joint-space waypoints.
Parameters
----------
space : ompl.base.StateSpace
State space used to allocate waypoint states.
si : ompl.base.SpaceInformation
OMPL space-information object associated with the path.
waypoints : array_like, shape (m, n)
Joint-space waypoints.
Returns
-------
ompl.geometric.PathGeometric
OMPL path containing the supplied waypoints in order.
Raises
------
ValueError
If any waypoint length does not match ``space.getDimension()``.
"""
path = og.PathGeometric(si)
for waypoint in waypoints:
_append_state(path, vec_to_state(space, waypoint))
return path
[docs]
class JointPathPlanner:
"""OMPL/MuJoCo joint-space path planner.
Attributes
----------
space : ompl.base.RealVectorStateSpace
OMPL joint-space state space.
ss : ompl.geometric.SimpleSetup
OMPL planning setup.
si : ompl.base.SpaceInformation
OMPL space-information object associated with ``ss``.
planner : ompl.base.Planner or None
Currently configured planner instance.
qaddr : list of int
MuJoCo ``qpos`` addresses for the planned joints.
bounds : ndarray, shape (n, 2)
Joint lower and upper bounds.
robot_geoms : set of int
Geom ids considered robot geoms during collision checking.
excluded_contact_geoms : set of int
Geom ids ignored during built-in MuJoCo contact checking.
excluded_contact_geom_pairs : set of tuple of int
Unordered geom-id pairs ignored during built-in MuJoCo contact checking.
included_contact_geoms : set of int
Geom ids whose contacts always invalidate a state.
has_welded_objects : list of int
Weld equality ids touching the robot body trees.
"""
[docs]
def __init__(
self,
robot: "robot",
q_start: VectorLike,
q_goal: VectorLike,
algorithm: Optional[str] = "RRTConnect",
max_planning_time: float = 5.0,
max_simplification_time: float = 0.0,
clearance: float = 0.0,
only_robot_contacts: bool = True,
excluded_contact_geom_ids: Optional[Iterable[int]] = None,
excluded_contact_geom_pairs: Optional[Iterable[ContactGeomPairLike]] = None,
included_contact_geoms: Optional[Iterable[int]] = None,
edge_resolution: float = 0.01,
interpolation_count: int = 0,
validity_fn: Optional[StateValidityFn] = None,
objective: Optional[ob.OptimizationObjective] = None,
) -> None:
"""Initialize the planner and configure the OMPL problem.
Parameters
----------
robot : robot-like
RobotBlockSet MuJoCo robot object. A composite object exposing
``robots`` is also supported; each child robot must expose ``scene``,
``JointNames`` and ``BaseName``.
q_start : array_like, shape (n,)
Start joint configuration in the order defined by the robot joint
names.
q_goal : array_like, shape (n,)
Goal joint configuration in the same joint order as ``q_start``.
algorithm : str or None, optional
OMPL planner name. ``"RRTConnect"`` is used by default. Supported
names include ``"PRM*"``, ``"LazyPRM*"``, ``"RRT*"``, ``"RRT#"``,
``"RRTX"``, ``"AORRTC"``, ``"InformedRRT*"``, ``"BIT*"``,
``"BLIT*"``, ``"ABIT*"``, ``"AIT*"``, ``"EIT*"``, ``"LBTRRT"``,
``"SST"``, ``"T-RRT"``, ``"TRRT*"``, ``"ATRRT"``, ``"SPARS"``,
``"SPARS2"``, ``"FMT*"``, ``"ST-RRT*"``, ``"CForest"`` and
``"APS"``. If ``None``, OMPL's default planner selection is used.
max_planning_time : float, optional
Default time budget in seconds used by :meth:`plan`.
max_simplification_time : float, optional
Time budget in seconds passed to ``SimpleSetup.simplifySolution``.
The default ``0.0`` uses OMPL's default simplification procedure.
clearance : float, optional
Collision-margin inflation applied to robot geoms while planning and
validating paths.
only_robot_contacts : bool, optional
If ``True``, contacts not involving robot geoms are ignored.
excluded_contact_geom_ids : iterable of int, optional
MuJoCo geom ids whose contacts are ignored by the default MuJoCo
validity checker. If a contact involves any excluded geom, that contact
does not invalidate the state. This argument is not applied when a
custom ``validity_fn`` is supplied.
excluded_contact_geom_pairs : iterable of pair-like, optional
MuJoCo geom-id pairs whose contacts are ignored by the default MuJoCo
validity checker. Pair order is ignored, so ``(g1, g2)`` and
``(g2, g1)`` are equivalent. This argument is not applied when a custom
``validity_fn`` is supplied.
included_contact_geoms : iterable of int, optional
MuJoCo geom ids whose contacts always invalidate the state in the
default MuJoCo validity checker. If any contact involves one of these
geoms, the state is invalid even if that contact would otherwise be
ignored by ``only_robot_contacts`` or the exclusion filters. This
argument is not applied when a custom ``validity_fn`` is supplied.
edge_resolution : float, optional
OMPL validity-checking resolution as a fraction of state-space extent.
interpolation_count : int, optional
Number of path states after interpolation. Values greater than ``3``
enable interpolation; values ``0`` through ``3`` leave the simplified
solution path unchanged.
validity_fn : callable, optional
Custom OMPL state-validity callback. If omitted, MuJoCo contact
checking is used.
objective : ompl.base.OptimizationObjective, optional
Optimization objective for planners that require one. If omitted, a
path-length objective is created for optimizing planners.
Returns
-------
None
Raises
------
ValueError
If ``q_start``/``q_goal`` dimensions do not match the robot joints, if
``max_simplification_time`` is negative, if an excluded geom pair does
not contain exactly two ids, or if the planner name is not supported by
the installed OMPL build.
"""
self.robot = robot
self.q_start = np.asarray(q_start, dtype=float)
self.q_goal = np.asarray(q_goal, dtype=float)
self.algorithm = algorithm
self.max_planning_time = float(max_planning_time)
self.max_simplification_time = float(max_simplification_time)
if self.max_simplification_time < 0.0:
raise ValueError("max_simplification_time must be non-negative")
self.clearance = float(clearance)
self.only_robot_contacts = bool(only_robot_contacts)
self.excluded_contact_geoms = set() if excluded_contact_geom_ids is None else {int(g) for g in excluded_contact_geom_ids}
self.included_contact_geoms = set() if included_contact_geoms is None else {int(g) for g in included_contact_geoms}
self.excluded_contact_geom_pairs: set[tuple[int, int]] = set()
if excluded_contact_geom_pairs is not None:
for pair in excluded_contact_geom_pairs:
pair_tuple = tuple(pair)
if len(pair_tuple) != 2:
raise ValueError("Each excluded contact geom pair must contain exactly two geom ids")
geom1, geom2 = int(pair_tuple[0]), int(pair_tuple[1])
self.excluded_contact_geom_pairs.add((geom1, geom2) if geom1 <= geom2 else (geom2, geom1))
self.edge_resolution = float(edge_resolution)
self.interpolation_count = int(interpolation_count)
self.custom_validity_fn = validity_fn
self.objective = objective
self._resolve_robot_model()
self.qaddr, self.bounds = get_robot_joints_data(self.model, self.joint_names)
self.n = int(self.bounds.shape[0])
if len(self.q_start) != self.n or len(self.q_goal) != self.n:
raise ValueError(f"Expected q_start/q_goal of length {self.n}; got {len(self.q_start)} and {len(self.q_goal)}")
self.robot_bodies_ids = get_body_descendants(self.model, self.base_names, include_self=True, del_dup=True)
self.robot_geoms_ids = get_geoms_of_body(self.model, self.robot_bodies_ids)
self.robot_geoms: set[int] = {int(g) for g in self.robot_geoms_ids}
self.has_welded_objects: list[int] = []
seen_welds: set[int] = set()
for base_name in self.base_names:
for weld_id in check_object_weld(self.model, self.data, base_name):
if weld_id not in seen_welds:
self.has_welded_objects.append(weld_id)
seen_welds.add(weld_id)
if self.only_robot_contacts:
for base_name in self.base_names:
self.included_contact_geoms.update(geoms_of_welded_objects(self.model, self.data, base_name))
self.qpos_template = np.array(self.data.qpos, copy=True)
self._data = copy(self.data) # Used in state validity fcn
self.space = make_state_space(self.bounds)
self.ss = og.SimpleSetup(self.space)
self.validity_fn = self.custom_validity_fn if self.custom_validity_fn is not None else self._is_state_valid
_set_state_validity_checker(self.ss, self.validity_fn)
self.si = self.ss.getSpaceInformation()
self.si.setStateValidityCheckingResolution(self.edge_resolution)
self.planner: Optional[ob.Planner] = None
self.set_planner(self.algorithm, objective=self.objective)
self.set_start_goal(self.q_start, self.q_goal)
def _resolve_robot_model(self) -> None:
"""Resolve robot scene, model, data and joint/body names.
Returns
-------
None
Notes
-----
Composite robot objects are detected by the presence of a ``robots``
attribute. The first child robot supplies the shared MuJoCo scene.
"""
if hasattr(self.robot, "robots"):
robots = self.robot.robots
self.scene = robots[0].scene
self.model = self.scene.model
self.data = self.scene.data
self.joint_names = [name for rbt in robots for name in rbt.JointNames]
self.base_names: Iterable[str] = [rbt.BaseName for rbt in robots]
else:
self.scene = self.robot.scene
self.model = self.scene.model
self.data = self.scene.data
self.joint_names = list(self.robot.JointNames)
self.base_names = [self.robot.BaseName]
[docs]
def update_robot_model_and_scene_data(
self,
robot: Optional["robot"] = None,
*,
clear_solution: bool = False,
) -> None:
"""Refresh cached robot, MuJoCo model and scene data references.
Call this after recompiling or replacing the robot scene so collision
checking uses the current ``scene.model`` and ``scene.data``. Existing
OMPL planning data is preserved unless ``clear_solution`` is ``True``.
If joint bounds change, the OMPL setup is rebuilt and existing planning
data cannot be preserved.
Parameters
----------
robot : robot-like, optional
Replacement robot or composite robot object. If omitted, the
existing ``self.robot`` object is queried again.
clear_solution : bool, optional
If ``True``, clear existing OMPL solution/planning data after
refreshing the MuJoCo references. Defaults to ``False``.
Returns
-------
None
Raises
------
ValueError
If the refreshed model has a different joint dimension than the
stored start or goal configuration.
"""
if robot is not None:
self.robot = robot
previous_bounds = getattr(self, "bounds", None)
self._resolve_robot_model()
qaddr, bounds = get_robot_joints_data(self.model, self.joint_names)
n = int(bounds.shape[0])
if len(self.q_start) != n or len(self.q_goal) != n:
raise ValueError(f"Expected q_start/q_goal of length {n}; got {len(self.q_start)} and {len(self.q_goal)}")
self.qaddr = qaddr
self.bounds = bounds
self.n = n
self.robot_bodies_ids = get_body_descendants(self.model, self.base_names, include_self=True, del_dup=True)
self.robot_geoms_ids = get_geoms_of_body(self.model, self.robot_bodies_ids)
self.robot_geoms = {int(g) for g in self.robot_geoms_ids}
self.qpos_template = np.array(self.data.qpos, copy=True)
self._data = copy(self.data) # Used in state validity fcn
bounds_changed = previous_bounds is None or previous_bounds.shape != bounds.shape or not np.allclose(previous_bounds, bounds)
if bounds_changed:
self.space = make_state_space(self.bounds)
self.ss = og.SimpleSetup(self.space)
self.validity_fn = self.custom_validity_fn if self.custom_validity_fn is not None else self._is_state_valid
_set_state_validity_checker(self.ss, self.validity_fn)
self.si = self.ss.getSpaceInformation()
self.si.setStateValidityCheckingResolution(self.edge_resolution)
algorithm = self.algorithm
objective = self.objective
self.planner = None
self.set_planner(algorithm, objective=objective)
self.set_start_goal(self.q_start, self.q_goal)
return
self.validity_fn = self.custom_validity_fn if self.custom_validity_fn is not None else self._is_state_valid
_set_state_validity_checker(self.ss, self.validity_fn)
self.si.setStateValidityCheckingResolution(self.edge_resolution)
if clear_solution:
if hasattr(self.ss, "clear"):
self.ss.clear()
else:
pdef = self.ss.getProblemDefinition()
if hasattr(pdef, "clearSolutionPaths"):
pdef.clearSolutionPaths()
def _is_state_valid(self, state: ob.State) -> bool:
"""Check whether an OMPL state is collision-free in MuJoCo.
Parameters
----------
state : ompl.base.State
Candidate joint-space state.
Returns
-------
bool
``True`` if the state is collision-free under the configured
contact-filtering policy, otherwise ``False``.
"""
q = self.qpos_template.copy()
for i, adr in enumerate(self.qaddr):
q[adr] = float(state[i])
if not self.has_welded_objects:
self._data.qpos[:] = q
self._data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self._data)
else:
for weld_id in self.has_welded_objects:
self._data.eq_active[weld_id] = 0
self._data.qpos[:] = q
self._data.qvel[:] = 0.0
mujoco.mj_forward(self.model, self._data)
for weld_id in self.has_welded_objects:
self._data.eq_active[weld_id] = 1
for base_name in self.base_names:
weld_all_objects(self.model, self._data, base_name)
mujoco.mj_forward(self.model, self._data)
for k in range(self._data.ncon):
contact = self._data.contact[k]
geom1, geom2 = int(contact.geom1), int(contact.geom2)
if self.included_contact_geoms and (geom1 in self.included_contact_geoms or geom2 in self.included_contact_geoms):
return False
if geom1 in self.excluded_contact_geoms or geom2 in self.excluded_contact_geoms:
continue
if self.excluded_contact_geom_pairs:
contact_pair = (geom1, geom2) if geom1 <= geom2 else (geom2, geom1)
if contact_pair in self.excluded_contact_geom_pairs:
continue
if self.only_robot_contacts and (geom1 not in self.robot_geoms and geom2 not in self.robot_geoms):
continue
return False
return True
[docs]
def set_planner(
self,
algorithm: Optional[str] = None,
*,
objective: Optional[ob.OptimizationObjective] = None,
) -> Optional[ob.Planner]:
"""Configure and attach an OMPL planner.
Parameters
----------
algorithm : str or None, optional
Planner name. If ``None``, OMPL's default planner selection is left
in place.
objective : ompl.base.OptimizationObjective, optional
Optimization objective used by optimizing planners. If omitted, the
current ``self.objective`` is reused; if that is also ``None``, a
path-length objective is created when needed.
Returns
-------
ompl.base.Planner or None
Configured planner instance, or ``None`` when ``algorithm`` is
``None``.
Raises
------
ValueError
If the planner name is unknown or unavailable in the installed OMPL
build.
"""
self.algorithm = algorithm
if objective is None:
objective = self.objective
self.objective = objective
if algorithm is None:
self.planner = None
elif check_option(algorithm, "RRTConnect"):
self.planner = og.RRTConnect(self.si)
self.ss.setPlanner(self.planner)
else:
if objective is None:
objective = ob.PathLengthOptimizationObjective(self.si)
self.objective = objective
self.ss.setOptimizationObjective(objective)
key = algorithm.strip().lower().replace(" ", "")
key = key.replace("star", "*").replace("rrtsharp", "rrt#").replace("strrt*", "st-rrt*")
planner_table = {
"aorrtc": getattr(og, "AORRTC", None), # AORRTC
"prm*": getattr(og, "PRMstar", None), # PRM*
"lazyprm*": getattr(og, "LazyPRMstar", None), # LazyPRM*
"rrt*": getattr(og, "RRTstar", None), # RRT*
"rrt#": getattr(og, "RRTsharp", None), # RRT#
"rrtx": getattr(og, "RRTXstatic", None), # RRTX (static variant)
"informedrrt*": getattr(og, "InformedRRTstar", None), # Informed RRT*
"bit*": getattr(og, "BITstar", None), # BIT*
"blit*": getattr(og, "BLITstar", None), # BLIT*
"abit*": getattr(og, "ABITstar", None), # ABIT*
"ait*": getattr(og, "AITstar", None), # AIT*
"eit*": getattr(og, "EITstar", None), # EIT*
"lbtrrt": getattr(og, "LBTRRT", None), # LBTRRT
"sst": getattr(og, "SST", None), # Sparse Stable RRT
"t-rrt": getattr(og, "TRRT", None), # T-RRT
"trrt*": getattr(og, "TRRTstar", None), # TRRT*
"t-rrt*": getattr(og, "TRRTstar", None), # T-RRT*
"atrrt": getattr(og, "ATRRT", None), # ATRRT
"at-rrt": getattr(og, "ATRRT", None), # AT-RRT
"spars": getattr(og, "SPARS", None), # SPARS
"spars2": getattr(og, "SPARStwo", None), # SPARS2
"fmt*": getattr(og, "FMT", None), # FMT*
"st-rrt*": getattr(og, "STRRTstar", None), # ST-RRT*
}
if key in ("cforest", "aps", "anytimepathshortening"):
if key == "cforest":
PlannerCls = getattr(og, "CForest", None)
subplanners = {
"RRTstar": getattr(og, "RRTstar", None),
"BITstar": getattr(og, "BITstar", None),
}
else:
PlannerCls = getattr(og, "AnytimePathShortening", None)
subplanners = {
"BITstar": getattr(og, "BITstar", None),
"InformedRRTstar": getattr(og, "InformedRRTstar", None),
}
if PlannerCls is None:
raise ValueError(f"Planner '{algorithm}' not available in this OMPL build")
missing = [name for name, cls in subplanners.items() if cls is None]
if missing:
raise ValueError(f"Planner '{algorithm}' requires unavailable OMPL planner(s): {', '.join(missing)}")
self.planner = PlannerCls(self.si)
for PlannerSubCls in subplanners.values():
self.planner.addPlanner(PlannerSubCls(self.si))
self.ss.setPlanner(self.planner)
else:
if key not in planner_table:
raise ValueError(f"Unknown/unsupported planner '{algorithm}'")
PlannerCls = planner_table[key]
if PlannerCls is None:
raise ValueError(f"Planner '{algorithm}' not available in this OMPL build")
try:
self.planner = PlannerCls(self.si)
except Exception as e:
raise ValueError(f"Planner '{algorithm}' not available in this OMPL build") from e
self.ss.setPlanner(self.planner)
return self.planner
[docs]
def set_start_goal(
self,
q_start: VectorLike,
q_goal: VectorLike,
threshold: float = 1e-3,
) -> None:
"""Set the OMPL start and goal states.
Parameters
----------
q_start : array_like, shape (n,)
Start joint configuration.
q_goal : array_like, shape (n,)
Goal joint configuration.
threshold : float, optional
OMPL goal-threshold distance.
Returns
-------
None
Raises
------
ValueError
If ``q_start`` or ``q_goal`` has the wrong dimension.
"""
self.q_start = np.asarray(q_start, dtype=float)
self.q_goal = np.asarray(q_goal, dtype=float)
if len(self.q_start) != self.n or len(self.q_goal) != self.n:
raise ValueError(f"Expected q_start/q_goal of length {self.n}; got {len(self.q_start)} and {len(self.q_goal)}")
start_state = vec_to_state(self.space, self.q_start)
goal_state = vec_to_state(self.space, self.q_goal)
self.ss.setStartAndGoalStates(start_state, goal_state, threshold)
def _inflate_collision_margins(self) -> CollisionMarginState:
"""Temporarily inflate robot geom collision margins.
Returns
-------
tuple or None
Opaque margin state used by :meth:`_restore_collision_margins`, or
``None`` when no inflation is requested.
"""
if self.clearance <= 0.0:
return None
if hasattr(self.model, "geom_margin"):
idx = np.asarray(self.robot_geoms_ids, dtype=int)
prev_margins = self.model.geom_margin[idx].copy()
self.model.geom_margin[idx] = self.clearance
return idx, prev_margins
prev_margins = {g: self.scene.model.geom(g).margin for g in self.robot_geoms_ids}
for geom_id in self.robot_geoms_ids:
self.scene.model.geom(geom_id).margin = self.clearance
return None, prev_margins
def _restore_collision_margins(self, margin_state: CollisionMarginState) -> None:
"""Restore collision margins saved by :meth:`_inflate_collision_margins`.
Parameters
----------
margin_state : tuple or None
State returned by :meth:`_inflate_collision_margins`.
Returns
-------
None
"""
if margin_state is None:
return
idx, prev_margins = margin_state
if idx is not None and isinstance(prev_margins, np.ndarray):
self.model.geom_margin[idx] = prev_margins
else:
for geom_id, value in prev_margins.items():
self.scene.model.geom(geom_id).margin = float(value)
def _solution_arrays(self) -> Tuple[np.ndarray, np.ndarray]:
"""Convert the current OMPL solution to NumPy arrays.
Returns
-------
path : ndarray, shape (m, n)
Simplified solution path. If ``self.interpolation_count > 3``, the
path is interpolated to that many states before conversion.
waypoints : ndarray, shape (k, n)
Raw OMPL solution states before simplification/interpolation.
"""
path = self.ss.getSolutionPath()
waypoints = [state_to_vec(path.getState(i), self.n) for i in range(path.getStateCount())]
self.ss.simplifySolution(self.max_simplification_time)
path = self.ss.getSolutionPath()
if self.interpolation_count > 3:
path.interpolate(self.interpolation_count)
path_int = [state_to_vec(path.getState(i), self.n) for i in range(path.getStateCount())]
return np.asarray(path_int, dtype=float).reshape(-1, self.n), np.asarray(waypoints, dtype=float).reshape(-1, self.n)
def _has_solution(self) -> bool:
"""Return whether the current OMPL problem has a solution.
Returns
-------
bool
``True`` if the problem definition reports a solution or a solution
path can be retrieved, otherwise ``False``.
"""
pdef = self.ss.getProblemDefinition()
if hasattr(pdef, "hasSolution"):
return bool(pdef.hasSolution())
try:
self.ss.getSolutionPath()
except Exception:
return False
return True
[docs]
def is_path_valid(
self,
path: Optional[PathLike] = None,
refresh_scene_state: bool = True,
) -> bool:
"""Check whether a path is valid in the current scene.
If ``path`` is omitted, the current OMPL solution path is checked. If a
waypoint array is supplied, it is converted to an OMPL ``PathGeometric``
and checked with the planner's current validity checker and edge
resolution.
Parameters
----------
path : array_like, shape (m, n), optional
Joint-space path to validate. If omitted, the current OMPL solution
path is validated.
refresh_scene_state : bool, optional
If ``True``, refresh the nominal MuJoCo ``qpos`` template from
``self.data.qpos`` before checking.
Returns
-------
bool
``True`` when all states and motions in the path are valid.
Raises
------
ValueError
If ``path`` has an invalid shape or contains non-finite values.
"""
if refresh_scene_state:
self.qpos_template = np.array(self.data.qpos, copy=True)
margin_state = self._inflate_collision_margins()
try:
if hasattr(self.si, "isSetup") and not self.si.isSetup():
self.si.setup()
if path is None:
if not self._has_solution():
return False
return bool(self.ss.getSolutionPath().check())
waypoints = np.asarray(path, dtype=float)
if waypoints.ndim == 1:
waypoints = waypoints.reshape(1, -1)
if waypoints.ndim != 2 or waypoints.shape[1] != self.n:
raise ValueError(f"Expected path with shape (m, {self.n}); got {waypoints.shape}")
if not np.all(np.isfinite(waypoints)):
raise ValueError("path contains non-finite values")
path_geo = _path_geometric_from_waypoints(self.space, self.si, waypoints)
return bool(path_geo.check())
finally:
self._restore_collision_margins(margin_state)
[docs]
def valid_path_segments(
self,
path: Optional[PathLike] = None,
*,
refresh_scene_state: bool = True,
) -> list[tuple[int, int]]:
"""Return maximal valid waypoint-index segments for a path.
Each returned pair is ``(first_index, last_index)``. A segment is valid
when all included waypoint states are valid and all motions between
consecutive included waypoints are valid.
Parameters
----------
path : array_like, shape (m, n), optional
Joint-space path to segment. If omitted, the current OMPL solution
path is used.
refresh_scene_state : bool, optional
If ``True``, refresh the nominal MuJoCo ``qpos`` template from
``self.data.qpos`` before checking.
Returns
-------
list of tuple of int
Maximal valid segments as inclusive waypoint-index pairs.
Raises
------
ValueError
If ``path`` has an invalid shape or contains non-finite values.
"""
if refresh_scene_state:
self.qpos_template = np.array(self.data.qpos, copy=True)
if path is None:
if not self._has_solution():
return []
solution = self.ss.getSolutionPath()
waypoints = np.asarray([state_to_vec(solution.getState(i), self.n) for i in range(solution.getStateCount())], dtype=float).reshape(-1, self.n)
else:
waypoints = np.asarray(path, dtype=float)
if waypoints.ndim == 1:
waypoints = waypoints.reshape(1, -1)
if waypoints.ndim != 2 or waypoints.shape[1] != self.n:
raise ValueError(f"Expected path with shape (m, {self.n}); got {waypoints.shape}")
if not np.all(np.isfinite(waypoints)):
raise ValueError("path contains non-finite values")
if waypoints.shape[0] == 0:
return []
margin_state = self._inflate_collision_margins()
try:
if hasattr(self.si, "isSetup") and not self.si.isSetup():
self.si.setup()
states = [_as_ompl_state(vec_to_state(self.space, waypoint)) for waypoint in waypoints]
state_valid = [bool(self.si.isValid(state)) for state in states]
segments: list[tuple[int, int]] = []
start: Optional[int] = None
last: Optional[int] = None
for i, is_valid in enumerate(state_valid):
if not is_valid:
if start is not None and last is not None:
segments.append((start, last))
start = None
last = None
continue
if start is None:
start = i
last = i
continue
if bool(self.si.checkMotion(states[i - 1], states[i])):
last = i
else:
if last is not None:
segments.append((start, last))
start = i
last = i
if start is not None and last is not None:
segments.append((start, last))
return segments
finally:
self._restore_collision_margins(margin_state)
[docs]
def plan(
self,
q_start: Optional[VectorLike] = None,
q_goal: Optional[VectorLike] = None,
max_planning_time: Optional[float] = None,
) -> PathResult:
"""Solve the configured planning query.
Parameters
----------
q_start : array_like, shape (n,), optional
Replacement start configuration. If omitted, ``self.q_start`` is
used.
q_goal : array_like, shape (n,), optional
Replacement goal configuration. If omitted, ``self.q_goal`` is used.
max_planning_time : float, optional
Solve time budget in seconds. If omitted, ``self.max_planning_time``
is used.
Returns
-------
path : ndarray, shape (m, n), or None
Simplified and optionally interpolated solution path.
waypoints : ndarray, shape (k, n), or None
Raw solution waypoints before simplification.
Raises
------
ValueError
If replacement start or goal configurations have invalid
dimensions.
"""
if q_start is not None or q_goal is not None:
self.set_start_goal(
self.q_start if q_start is None else q_start,
self.q_goal if q_goal is None else q_goal,
)
allowed_time = self.max_planning_time if max_planning_time is None else float(max_planning_time)
margin_state = self._inflate_collision_margins()
try:
if self.ss.solve(allowed_time):
return self._solution_arrays()
if self._has_solution():
return self._solution_arrays()
return None, None
finally:
self._restore_collision_margins(margin_state)
[docs]
def initialize_path(
self,
path: PathLike,
*,
validate_path: bool = True,
clear_existing: bool = True,
set_start_goal_from_path: bool = True,
threshold: float = 1e-3,
name: str = "initial_path",
refresh_scene_state: bool = True,
) -> PathResult:
"""Initialize the current OMPL solution from path waypoints.
This method seeds the planner with an existing collision-free path so
that :meth:`refine_solution` or :meth:`refine_with_planner` can improve
it as if it had been found by :meth:`plan`.
Parameters
----------
path : array_like, shape (m, n)
Joint-space waypoint path to insert as the incumbent solution.
validate_path : bool, optional
If ``True``, verify that all states and motions in ``path`` are
valid before adding it to OMPL.
clear_existing : bool, optional
If ``True``, clear existing OMPL planning data before adding the
supplied path.
set_start_goal_from_path : bool, optional
If ``True``, set the planning query start and goal to the first and
last path waypoints.
threshold : float, optional
Goal threshold used when ``set_start_goal_from_path`` is ``True``.
name : str, optional
Name attached to the inserted OMPL solution path.
refresh_scene_state : bool, optional
If ``True`` and ``validate_path`` is also ``True``, refresh the
nominal MuJoCo ``qpos`` template from ``self.data.qpos`` before
validating the path.
Returns
-------
initialized_path : ndarray, shape (m, n)
Inserted path as a floating-point NumPy array.
initialized_waypoints : ndarray, shape (m, n)
Copy of ``initialized_path`` for compatibility with :meth:`plan`.
Raises
------
ValueError
If ``path`` has an invalid shape, contains non-finite values, has
fewer than two waypoints, fails validation, or has endpoints that do
not match the current query when ``set_start_goal_from_path`` is
``False``.
"""
waypoints = np.asarray(path, dtype=float)
if waypoints.ndim == 1:
waypoints = waypoints.reshape(1, -1)
if waypoints.ndim != 2 or waypoints.shape[1] != self.n:
raise ValueError(f"Expected path with shape (m, {self.n}); got {waypoints.shape}")
if waypoints.shape[0] < 2:
raise ValueError("path must contain at least two waypoints")
if not np.all(np.isfinite(waypoints)):
raise ValueError("path contains non-finite values")
if set_start_goal_from_path:
self.set_start_goal(waypoints[0], waypoints[-1], threshold=threshold)
elif not (np.allclose(waypoints[0], self.q_start) and np.allclose(waypoints[-1], self.q_goal)):
raise ValueError("path endpoints do not match the current start and goal")
if validate_path and not self.is_path_valid(waypoints, refresh_scene_state=refresh_scene_state):
raise ValueError("path is not collision-free according to the configured validity checker")
if clear_existing:
if hasattr(self.ss, "clear"):
self.ss.clear()
else:
pdef = self.ss.getProblemDefinition()
if hasattr(pdef, "clearSolutionPaths"):
pdef.clearSolutionPaths()
path_geo = _path_geometric_from_waypoints(self.space, self.si, waypoints)
pdef = self.ss.getProblemDefinition()
pdef.addSolutionPath(path_geo, False, 0.0, name)
initialized_path = waypoints.copy()
return initialized_path, initialized_path.copy()
[docs]
def refine_solution(self, additional_time: float) -> PathResult:
"""Continue solving the current planning problem.
Repeated calls to ``SimpleSetup.solve`` let optimizing planners improve
the incumbent solution without rebuilding the planning problem. Call
:meth:`plan` or :meth:`initialize_path` first to create an incumbent.
Parameters
----------
additional_time : float
Extra solve time budget in seconds. A value of ``0`` returns the
current solution if one exists.
Returns
-------
path : ndarray, shape (m, n), or None
Simplified and optionally interpolated solution path.
waypoints : ndarray, shape (k, n), or None
Raw solution waypoints before simplification.
Raises
------
ValueError
If ``additional_time`` is negative.
"""
additional_time = float(additional_time)
if additional_time < 0.0:
raise ValueError("additional_time must be non-negative")
margin_state = self._inflate_collision_margins()
try:
if additional_time > 0.0:
self.ss.solve(additional_time)
if self._has_solution():
return self._solution_arrays()
return None, None
finally:
self._restore_collision_margins(margin_state)
[docs]
def refine_with_planner(
self,
algorithm: str,
additional_time: float,
*,
objective: Optional[ob.OptimizationObjective] = None,
) -> PathResult:
"""Switch planners and continue from the current solution.
Call :meth:`plan` first with a fast planner such as ``RRTConnect`` or
seed an existing path with :meth:`initialize_path`, then use this method
with an optimizing planner such as ``InformedRRT*`` or ``BIT*`` to
improve the path without losing the existing solution.
Parameters
----------
algorithm : str
Planner name used for refinement.
additional_time : float
Extra solve time budget in seconds.
objective : ompl.base.OptimizationObjective, optional
Optimization objective for the refinement planner.
Returns
-------
path : ndarray, shape (m, n), or None
Simplified and optionally interpolated solution path.
waypoints : ndarray, shape (k, n), or None
Raw solution waypoints before simplification.
Raises
------
ValueError
If ``additional_time`` is negative or the requested planner is
unavailable.
RuntimeError
If no incumbent solution is available.
"""
additional_time = float(additional_time)
if additional_time < 0.0:
raise ValueError("additional_time must be non-negative")
if not self._has_solution():
raise RuntimeError("No solution available to refine; call plan() or initialize_path() first")
previous_algorithm = self.algorithm or "existing_solution"
incumbent_path = og.PathGeometric(self.ss.getSolutionPath())
self.set_planner(algorithm, objective=objective)
pdef = self.ss.getProblemDefinition()
if hasattr(pdef, "clearSolutionPaths"):
pdef.clearSolutionPaths()
pdef.addSolutionPath(incumbent_path, False, 0.0, f"seed_{previous_algorithm}")
margin_state = self._inflate_collision_margins()
try:
if additional_time > 0.0:
self.ss.solve(additional_time)
if self._has_solution():
return self._solution_arrays()
return None, None
finally:
self._restore_collision_margins(margin_state)
[docs]
def repair_path(
self,
path: Optional[PathLike] = None,
*,
algorithm: Optional[str] = None,
max_planning_time: Optional[float] = None,
max_simplification_time: Optional[float] = None,
interpolation_count: Optional[int] = None,
objective: Optional[ob.OptimizationObjective] = None,
refresh_scene_state: bool = True,
) -> PathResult:
"""Repair a path by replanning invalid spans between valid segments.
The method finds valid path segments, replans each gap between
consecutive valid segments, and concatenates the original valid segments
with the newly planned bridges. If the first or last waypoint is invalid,
the path cannot be repaired because there is no valid endpoint to anchor
the repair.
Parameters
----------
path : array_like, shape (m, n), optional
Path to repair. If omitted, the current OMPL solution path is used.
algorithm : str or None, optional
Planner used for bridge segments. If omitted, ``self.algorithm`` is
reused.
max_planning_time : float, optional
Bridge-planning time budget in seconds. If omitted,
``self.max_planning_time`` is used.
max_simplification_time : float, optional
Bridge simplification time budget in seconds. If omitted,
``self.max_simplification_time`` is used.
interpolation_count : int, optional
Bridge interpolation count. If omitted, ``self.interpolation_count``
is used.
objective : ompl.base.OptimizationObjective, optional
Optimization objective used by bridge planners.
refresh_scene_state : bool, optional
If ``True``, refresh the nominal MuJoCo ``qpos`` template from
``self.data.qpos`` before computing valid segments.
Returns
-------
repaired_path : ndarray, shape (r, n), or None
Repaired complete path, or ``None`` if repair fails.
repaired_waypoints : ndarray, shape (r, n), or None
Copy of ``repaired_path`` for compatibility with :meth:`plan`.
``None`` if repair fails.
Raises
------
ValueError
If ``path`` has an invalid shape, contains non-finite values, or a
supplied time budget is negative. A ``ValueError`` can also be
raised if a requested bridge planner is unavailable.
"""
if path is None:
if not self._has_solution():
return None, None
solution = self.ss.getSolutionPath()
waypoints = np.asarray([state_to_vec(solution.getState(i), self.n) for i in range(solution.getStateCount())], dtype=float).reshape(-1, self.n)
else:
waypoints = np.asarray(path, dtype=float)
if waypoints.ndim == 1:
waypoints = waypoints.reshape(1, -1)
if waypoints.ndim != 2 or waypoints.shape[1] != self.n:
raise ValueError(f"Expected path with shape (m, {self.n}); got {waypoints.shape}")
if not np.all(np.isfinite(waypoints)):
raise ValueError("path contains non-finite values")
if waypoints.shape[0] == 0:
return None, None
valid_segments = self.valid_path_segments(waypoints, refresh_scene_state=refresh_scene_state)
self.repair_valid_segments = valid_segments
self.repair_spans: list[tuple[int, int]] = []
if not valid_segments:
return None, None
if valid_segments[0][0] != 0 or valid_segments[-1][1] != waypoints.shape[0] - 1:
return None, None
if len(valid_segments) == 1 and valid_segments[0] == (0, waypoints.shape[0] - 1):
repaired = waypoints.copy()
return repaired, repaired.copy()
repair_algorithm = self.algorithm if algorithm is None else algorithm
repair_time = self.max_planning_time if max_planning_time is None else float(max_planning_time)
if repair_time < 0.0:
raise ValueError("max_planning_time must be non-negative")
repair_simplification_time = self.max_simplification_time if max_simplification_time is None else float(max_simplification_time)
if repair_simplification_time < 0.0:
raise ValueError("max_simplification_time must be non-negative")
repair_interpolation_count = self.interpolation_count if interpolation_count is None else int(interpolation_count)
path_parts: list[np.ndarray] = []
def append_path_part(part: np.ndarray) -> None:
"""Append a path part while avoiding duplicate boundary states.
Parameters
----------
part : array_like, shape (p, n)
Path segment to append.
Returns
-------
None
"""
part = np.asarray(part, dtype=float).reshape(-1, self.n)
if path_parts and part.shape[0] > 0 and np.allclose(path_parts[-1][-1], part[0]):
part = part[1:]
if part.shape[0] > 0:
path_parts.append(part)
first_start, first_end = valid_segments[0]
append_path_part(waypoints[first_start : first_end + 1])
for (_, prev_end), (next_start, next_end) in zip(valid_segments[:-1], valid_segments[1:]):
bridge_planner = JointPathPlanner(
self.robot,
waypoints[prev_end],
waypoints[next_start],
algorithm=repair_algorithm,
max_planning_time=repair_time,
max_simplification_time=repair_simplification_time,
clearance=self.clearance,
only_robot_contacts=self.only_robot_contacts,
excluded_contact_geom_ids=self.excluded_contact_geoms,
excluded_contact_geom_pairs=self.excluded_contact_geom_pairs,
edge_resolution=self.edge_resolution,
interpolation_count=repair_interpolation_count,
validity_fn=self.custom_validity_fn,
objective=objective,
included_contact_geoms=self.included_contact_geoms,
)
bridge_path, _ = bridge_planner.plan()
if bridge_path is None:
return None, None
if not bridge_planner.is_path_valid(bridge_path, refresh_scene_state=False):
return None, None
append_path_part(bridge_path)
append_path_part(waypoints[next_start : next_end + 1])
self.repair_spans.append((prev_end, next_start))
repaired = np.vstack(path_parts)
if not self.is_path_valid(repaired, refresh_scene_state=False):
return None, repaired
return repaired, repaired.copy()
[docs]
def plan_robot_motion(
robot: "robot",
q_start: VectorLike,
q_goal: VectorLike,
algorithm: Optional[str] = "RRTConnect",
max_planning_time: float = 5.0,
max_simplification_time: float = 0.0,
clearance: float = 0.0,
only_robot_contacts: bool = True,
excluded_contact_geom_ids: Optional[Iterable[int]] = None,
excluded_contact_geom_pairs: Optional[Iterable[ContactGeomPairLike]] = None,
edge_resolution: float = 0.01,
interpolation_count: int = 0,
validity_fn: Optional[StateValidityFn] = None,
objective: Optional[ob.OptimizationObjective] = None,
included_contact_geoms: Optional[Iterable[int]] = None,
) -> PathResult:
"""Plan a collision-free joint-space motion with OMPL and MuJoCo.
The function temporarily **inflates the robot geoms' collision margins** by
``clearance`` during planning (for a safety buffer) and restores the original
margins afterwards.
Parameters
----------
robot : robot-like
RobotBlockSet MuJoCo robot object, or a composite object exposing a
``robots`` sequence.
q_start : array_like, shape (n,)
Start joint configuration.
q_goal : array_like, shape (n,)
Goal joint configuration.
algorithm : str or None, optional
Planner name forwarded to :class:`JointPathPlanner`. If ``None``, OMPL's
default planner selection is used.
max_planning_time : float, optional
Planner solve time budget in seconds.
max_simplification_time : float, optional
Path-simplification time budget in seconds. The default ``0.0`` uses
OMPL's default simplification procedure.
clearance : float, optional
Collision-margin inflation applied to robot geoms while planning.
only_robot_contacts : bool, optional
If ``True``, contacts not involving robot geoms are ignored.
excluded_contact_geom_ids : iterable of int, optional
MuJoCo geom ids whose contacts are ignored by the default MuJoCo
validity checker. This argument is not applied when a custom
``validity_fn`` is supplied.
excluded_contact_geom_pairs : iterable of pair-like, optional
MuJoCo geom-id pairs whose contacts are ignored by the default MuJoCo
validity checker. Pair order is ignored. This argument is not applied
when a custom ``validity_fn`` is supplied.
edge_resolution : float, optional
OMPL validity-checking resolution as a fraction of state-space extent.
interpolation_count : int, optional
Number of path states after interpolation. Values greater than ``3``
enable interpolation; values ``0`` through ``3`` leave the simplified
path unchanged.
validity_fn : callable, optional
Custom OMPL state-validity callback. If omitted, MuJoCo contact
checking is used.
objective : ompl.base.OptimizationObjective, optional
Optimization objective for planners that require one.
included_contact_geoms : iterable of int, optional
MuJoCo geom ids whose contacts always invalidate the state in the
default MuJoCo validity checker. This argument is not applied when a
custom ``validity_fn`` is supplied.
Returns
-------
path : ndarray, shape (m, n), or None
Simplified and optionally interpolated solution path.
waypoints : ndarray, shape (k, n), or None
Raw solution waypoints before simplification.
Raises
------
ValueError
If joint dimensions do not match, ``max_simplification_time`` is
negative, an excluded geom pair does not contain exactly two ids, or
the requested planner is unavailable.
"""
planner = JointPathPlanner(
robot,
q_start,
q_goal,
algorithm=algorithm,
max_planning_time=max_planning_time,
max_simplification_time=max_simplification_time,
clearance=clearance,
only_robot_contacts=only_robot_contacts,
excluded_contact_geom_ids=excluded_contact_geom_ids,
excluded_contact_geom_pairs=excluded_contact_geom_pairs,
edge_resolution=edge_resolution,
validity_fn=validity_fn,
objective=objective,
interpolation_count=interpolation_count,
included_contact_geoms=included_contact_geoms,
)
return planner.plan()