pyroboplan 1.0.0__py3-none-any.whl

pyroboplan/__init__.py

@@ -0,0 +1 @@
+""" Educational Python library for manipulator motion planning. """

pyroboplan/core/__init__.py

@@ -0,0 +1,5 @@
+"""Core pyroboplan module.
+
+This module contains the essential utilities to represent a robot model for motion planning.
+It makes extensive use of the Pinocchio library but adds planning-specific abstractions over its implementation.
+"""

pyroboplan/core/planning_context.py

@@ -0,0 +1,18 @@
+import pinocchio
+
+
+class PlanningContext:
+    """
+    Defines a planning context, which contains Pinocchio models and necessary utilities for motion planning.
+    """
+
+    def __init__(self, model, visual_model=None, collision_model=None):
+        """
+        Creates a planning context instance given a set of Pinocchio models.
+        """
+        self.model = model
+        self.visual_model = visual_model
+        self.collision_model = collision_model
+
+        self.data = model.createData()
+        self.collision_data = collision_model.createData()

pyroboplan/core/utils.py

@@ -0,0 +1,358 @@
+""" Core utilities for robot modeling. """
+
+import copy
+import numpy as np
+import pinocchio
+
+
+def check_collisions_at_state(model, collision_model, q):
+    """
+    Checks whether a specified joint configuration is collision-free.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        q : array-like
+            The joint configuration of the model.
+
+    Returns
+    -------
+        bool
+            True is there are any collisions, otherwise False.
+    """
+    data = model.createData()
+    collision_data = collision_model.createData()
+    stop_at_first_collision = True  # For faster computation
+
+    pinocchio.computeCollisions(
+        model, data, collision_model, collision_data, q, stop_at_first_collision
+    )
+    return np.any([cr.isCollision() for cr in collision_data.collisionResults])
+
+
+def check_collisions_along_path(model, collision_model, q_path):
+    """
+    Checks whether a path consisting of multiple joint configurations is collision-free.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        q_path : list[array-like]
+            A list of joint configurations describing the path.
+
+    Returns
+    -------
+        bool
+            True is there are any collisions, otherwise False.
+    """
+    data = model.createData()
+    collision_data = collision_model.createData()
+    stop_at_first_collision = True  # For faster computation
+
+    for q in q_path:
+        pinocchio.computeCollisions(
+            model, data, collision_model, collision_data, q, stop_at_first_collision
+        )
+        if np.any([cr.isCollision() for cr in collision_data.collisionResults]):
+            return True
+
+    return False
+
+
+def configuration_distance(q_start, q_end):
+    """
+    Returns the distance between two joint configurations.
+
+    Parameters
+    ----------
+        q_start : array-like
+            The start joint configuration.
+        q_end : array-like
+            The end joint configuration.
+
+    Returns
+    -------
+        float
+            The distance between the two joint configurations.
+    """
+    return np.linalg.norm(q_end - q_start)
+
+
+def get_path_length(q_path):
+    """
+    Returns the configuration distance of a path.
+
+    Parameters
+    ----------
+        q_path : list[array-like]
+            A list of joint configurations describing a path.
+
+    Returns
+    -------
+        float
+            The total configuration distance of the entire path.
+    """
+    total_distance = 0.0
+    for idx in range(1, len(q_path)):
+        total_distance += configuration_distance(q_path[idx - 1], q_path[idx])
+    return total_distance
+
+
+def get_random_state(model, padding=0.0):
+    """
+    Returns a random state that is within the model's joint limits.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        padding : float or array-like, optional
+            The padding to use around the sampled joint limits.
+
+    Returns
+    -------
+        array-like
+            A set of randomly generated joint variables.
+    """
+    return np.random.uniform(
+        model.lowerPositionLimit + padding, model.upperPositionLimit - padding
+    )
+
+
+def get_random_collision_free_state(model, collision_model, padding=0.0, max_tries=100):
+    """
+    Returns a random state that is within the model's joint limits and is collision-free according to the collision model.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        padding : float or array-like, optional
+            The padding to use around the sampled joint limits.
+        max_tries : int, optional
+            The maximum number of tries for sampling a collision-free state.
+
+    Returns
+    -------
+        array-like
+            A set of randomly generated collision-free joint variables, or None if one cannot be found.
+    """
+    num_tries = 0
+    while num_tries < max_tries:
+        state = get_random_state(model, padding=padding)
+        if not check_collisions_at_state(model, collision_model, state):
+            return state
+        num_tries += 1
+
+    print(f"Could not generate collision-free state after {max_tries} tries.")
+    return None
+
+
+def get_random_transform(model, target_frame, joint_padding=0.0):
+    """
+    Returns a random transform for a target frame that is within the model's joint limits.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random transform.
+        target_frame : str
+            The name of the frame for which to generate a random transform.
+        joint_padding : float or array-like, optional
+            The padding to use around the sampled joint limits.
+
+    Returns
+    -------
+        `pinocchio.SE3`
+            A randomly generated transform for the specified target frame.
+    """
+    q_target = get_random_state(model, joint_padding)
+    data = model.createData()
+    target_frame_id = model.getFrameId(target_frame)
+    pinocchio.framesForwardKinematics(model, data, q_target)
+    return data.oMf[target_frame_id]
+
+
+def get_random_collision_free_transform(
+    model, collision_model, target_frame, joint_padding=0.0, max_tries=100
+):
+    """
+    Returns a random transform for a target frame that is within the model's joint limits and is collision-free.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random transform.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        target_frame : str
+            The name of the frame for which to generate a random transform.
+        joint_padding : float or array-like, optional
+            The padding to use around the sampled joint limits.
+        max_tries : int, optional
+            The maximum number of tries for sampling a collision-free state.
+
+    Returns
+    -------
+        `pinocchio.SE3`
+            A randomly generated transform for the specified target frame.
+    """
+    q_target = get_random_collision_free_state(
+        model, collision_model, joint_padding, max_tries
+    )
+    data = model.createData()
+    target_frame_id = model.getFrameId(target_frame)
+    pinocchio.framesForwardKinematics(model, data, q_target)
+    return data.oMf[target_frame_id]
+
+
+def check_within_limits(model, q):
+    """
+    Checks whether a particular joint configuration is within the model's joint limits.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to generate a random state.
+        q : array-like
+            The joint configuration for the model.
+
+    Returns
+    -------
+        bool
+            True if the configuration is within joint limits, otherwise False.
+    """
+    return np.all(q >= model.lowerPositionLimit) and np.all(
+        q <= model.upperPositionLimit
+    )
+
+
+def extract_cartesian_pose(model, target_frame, q, data=None):
+    """
+    Extracts the Cartesian pose of a specified model frame given a joint configuration.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to perform forward kinematics.
+        target_frame : str
+            The name of the target frame.
+        q : array-like
+            The joint configuration values describing the robot state.
+        data : `pinocchio.Data`, optional
+            The model data to use. If not set, one will be created.
+
+    Returns
+    -------
+        `pinocchio.SE3`
+            The transform describing the Cartesian pose of the specified frame at the provided joint configuration.
+    """
+    if data is None:
+        data = model.createData()
+    target_frame_id = model.getFrameId(target_frame)
+    pinocchio.framesForwardKinematics(model, data, q)
+    return copy.deepcopy(data.oMf[target_frame_id])
+
+
+def extract_cartesian_poses(model, target_frame, q_vec, data=None):
+    """
+    Extracts the Cartesian poses of a specified model frame given a list of joint configurations.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model from which to perform forward kinematics.
+        target_frame : str
+            The name of the target frame.
+        q_vec : array-like
+            A list of joint configuration values describing the path.
+        data : `pinocchio.Data`, optional
+            The model data to use. If not set, one will be created.
+
+    Returns
+    -------
+        list[`pinocchio.SE3`]
+            A list of transforms describing the Cartesian poses of the specified frame at the provided joint configurations.
+    """
+    if data is None:
+        data = model.createData()
+    target_frame_id = model.getFrameId(target_frame)
+    tforms = []
+    for q in q_vec:
+        pinocchio.framesForwardKinematics(model, data, q)
+        tforms.append(copy.deepcopy(data.oMf[target_frame_id]))
+    return tforms
+
+
+def get_collision_geometry_ids(model, collision_model, body):
+    """
+    Gets a list of collision geometry model IDs for a specified body name.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model to use for getting frame IDs.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        body : str
+            The name of the body.
+            This can be directly the name of a geometry in the collision model,
+            or it can be the name of the frame in the main model.
+
+    Return
+    ------
+        list[int]
+            A list of collision geometry IDs corresponding to the body name.
+    """
+    body_collision_geom_ids = []
+
+    # First, check if this is directly a collision geometry.
+    body_collision_geom_id = collision_model.getGeometryId(body)
+    if body_collision_geom_id < collision_model.ngeoms:
+        body_collision_geom_ids.append(body_collision_geom_id)
+    else:
+        # Otherwise, look for the frame name in the model and return its associated collision objects.
+        body_frame_id = model.getFrameId(body)
+        if body_frame_id < model.nframes:
+            for id, obj in enumerate(collision_model.geometryObjects):
+                if obj.parentFrame == body_frame_id:
+                    body_collision_geom_ids.append(id)
+
+    return body_collision_geom_ids
+
+
+def set_collisions(model, collision_model, body1, body2, enable):
+    """
+    Sets collision checking between two bodies by searching for their corresponding geometry objects in the collision model.
+
+    Parameters
+    ----------
+        model : `pinocchio.Model`
+            The model to use for getting frame IDs.
+        collision_model : `pinocchio.Model`
+            The model to use for collision checking.
+        body1 : str
+            The name of the first body.
+        body2 : str
+            The name of the second body.
+        enable : bool
+            If True, enables collisions. If False, disables collisions.
+    """
+    body1_collision_ids = get_collision_geometry_ids(model, collision_model, body1)
+    body2_collision_ids = get_collision_geometry_ids(model, collision_model, body2)
+
+    for id1 in body1_collision_ids:
+        for id2 in body2_collision_ids:
+            pair = pinocchio.CollisionPair(id1, id2)
+            if enable:
+                collision_model.addCollisionPair(pair)
+            else:
+                collision_model.removeCollisionPair(pair)

pyroboplan/ik/__init__.py

@@ -0,0 +1,4 @@
+""" Inverse kinematics module.
+
+This module contains solvers and utilities for inverse kinematics (IK).
+"""

pyroboplan/ik/differential_ik.py

@@ -0,0 +1,253 @@
+""" Utilities for differential IK. """
+
+import numpy as np
+import pinocchio
+import time
+
+from ..core.utils import (
+    check_collisions_at_state,
+    check_within_limits,
+    get_random_state,
+)
+from ..visualization.meshcat_utils import visualize_frame
+
+VIZ_INITIAL_RENDER_TIME = 0.5
+VIZ_SLEEP_TIME = 0.05
+
+
+class DifferentialIkOptions:
+    """Options for differential IK."""
+
+    def __init__(
+        self,
+        max_iters=200,
+        max_retries=10,
+        max_translation_error=1e-3,
+        max_rotation_error=1e-3,
+        damping=1e-3,
+        min_step_size=0.1,
+        max_step_size=0.5,
+    ):
+        """
+        Initializes a set of differential IK options.
+
+        Parameters
+        ----------
+            max_iters : int
+                Maximum number of iterations per try.
+            max_retries : int
+                Maximum number of retries with random restarts.
+                If set to 0, only the initial state provided will be used.
+            max_translation_error : float
+                Maximum translation error, in meters, to consider IK solved.
+            max_rotation_error : float
+                Maximum rotation error, in radians, to consider IK solved.
+            damping : float
+                Damping value, between 0 and 1, for the Jacobian pseudoinverse.
+                Setting this to a nonzero value is using Levenberg-Marquardt.
+            min_step_size : float
+                Minimum gradient step, between 0 and 1, based on ratio of current distance to target to initial distance to target.
+                To use a fixed step size, set both minimum and maximum values to be equal.
+            max_step_size : float
+                Maximum gradient step, between 0 and 1, based on ratio of current distance to target to initial distance to target.
+                To use a fixed step size, set both minimum and maximum values to be equal.
+        """
+        self.max_iters = max_iters
+        self.max_retries = max_retries
+        self.max_translation_error = max_translation_error
+        self.max_rotation_error = max_rotation_error
+        self.damping = damping
+        self.min_step_size = min_step_size
+        self.max_step_size = max_step_size
+
+
+class DifferentialIk:
+    """Differential IK solver.
+
+    This is a numerical IK solver that uses the manipulator's Jacobian to take first-order steps towards a solution.
+    It contains several of the common options such as damped least squares (Levenberg-Marquardt), random restarts, and nullspace projection.
+
+    Some good resources:
+      * https://motion.cs.illinois.edu/RoboticSystems/InverseKinematics.html
+      * https://homes.cs.washington.edu/~todorov/courses/cseP590/06_JacobianMethods.pdf
+      * https://www.cs.cmu.edu/~15464-s13/lectures/lecture6/iksurvey.pdf
+    """
+
+    def __init__(
+        self,
+        model,
+        collision_model=None,
+        data=None,
+        visualizer=None,
+        options=DifferentialIkOptions(),
+    ):
+        """
+        Creates an instance of a DifferentialIk solver.
+
+        Parameters
+        ----------
+            model : `pinocchio.Model`
+                The model to use for this solver.
+            collision_model : `pinocchio.Model`, optional
+                The model to use for collision checking. If None, no collision checking takes place.
+            data : `pinocchio.Data`, optional
+                The model data to use for this solver. If None, data is created automatically.
+            visualizer : `pinocchio.visualize.meshcat_visualizer.MeshcatVisualizer`, optional
+                The visualizer to use for this solver.
+            options : `DifferentialIkOptions`, optional
+                The options to use for solving IK. If not specified, default options are used.
+        """
+        self.model = model
+        self.collision_model = collision_model
+
+        if not data:
+            data = model.createData()
+        self.data = data
+
+        self.visualizer = visualizer
+        self.options = options
+
+    def solve(
+        self,
+        target_frame,
+        target_tform,
+        init_state=None,
+        nullspace_components=[],
+        verbose=False,
+    ):
+        """
+        Solves an IK query.
+
+        Parameters
+        ----------
+            target_frame : str
+                The name of the target frame in the model.
+            target_tform : `pinocchio.SE3`
+                The desired transformation of the target frame in the model.
+            init_state : array-like, optional
+                The initial state to solve from. If not specified, a random initial state will be selected.
+            nullspace_components : list[function], optional
+                An optional list of nullspace components to use when solving.
+                These components must take the form `lambda model, q: component(model, q, <other_args>)`.
+            verbose : bool, optional
+                If True, prints additional information to the console.
+
+        Returns
+        -------
+            array-like or None
+                A list of joint configuration values with the solution, if one was found. Otherwise, returns None.
+        """
+        target_frame_id = self.model.getFrameId(target_frame)
+
+        # Create a random initial state, if not specified
+        if init_state is None:
+            init_state = get_random_state(self.model)
+        if self.visualizer:
+            self.visualizer.displayFrames(True, frame_ids=[target_frame_id])
+            visualize_frame(self.visualizer, "ik_target_pose", target_tform)
+
+        # Initialize IK
+        solved = False
+        n_tries = 0
+        q_cur = init_state
+        initial_error_norm = None
+
+        if self.visualizer:
+            self.visualizer.display(q_cur)
+            time.sleep(VIZ_INITIAL_RENDER_TIME)  # Needed to render
+
+        while n_tries <= self.options.max_retries:
+            n_iters = 0
+            while n_iters < self.options.max_iters:
+                # Compute forward kinematics at the current state
+                pinocchio.framesForwardKinematics(self.model, self.data, q_cur)
+                cur_tform = self.data.oMf[target_frame_id]
+
+                # Check the error using actInv
+                error = target_tform.actInv(cur_tform)
+                error = -pinocchio.log(error).vector
+                if (
+                    np.linalg.norm(error[:3]) < self.options.max_translation_error
+                    and np.linalg.norm(error[3:]) < self.options.max_rotation_error
+                ):
+                    # Wrap to the range -/+ pi, and then check joint limits and collision.
+                    q_cur = (q_cur + np.pi) % (2 * np.pi) - np.pi
+                    if check_within_limits(self.model, q_cur):
+                        if self.collision_model is not None:
+                            if check_collisions_at_state(
+                                self.model, self.collision_model, q_cur
+                            ):
+                                if verbose:
+                                    print(
+                                        "Solved and within joint limits, but in collision."
+                                    )
+                            else:
+                                solved = True
+                                if verbose:
+                                    print(
+                                        "Solved, within joint limits, and collision-free!"
+                                    )
+                        else:
+                            solved = True
+                            if verbose:
+                                print("Solved and within joint limits!")
+                    else:
+                        if verbose:
+                            print("Solved, but outside joint limits.")
+                    break
+
+                # Calculate the Jacobian
+                J = pinocchio.computeFrameJacobian(
+                    self.model,
+                    self.data,
+                    q_cur,
+                    target_frame_id,
+                    pinocchio.ReferenceFrame.LOCAL,
+                )
+
+                # Solve for the gradient using damping and nullspace components,
+                # as specified
+                jjt = J.dot(J.T) + self.options.damping**2 * np.eye(6)
+
+                # Compute the gradient descent step size
+                error_norm = np.linalg.norm(error)
+                if not initial_error_norm:
+                    initial_error_norm = error_norm
+                alpha = self.options.min_step_size + (
+                    1.0 - error_norm / initial_error_norm
+                ) * (self.options.max_step_size - self.options.min_step_size)
+
+                # Gradient descent step
+                if not nullspace_components:
+                    q_cur += alpha * J.T @ np.linalg.solve(jjt, error)
+                else:
+                    nullspace_term = sum(
+                        [comp(self.model, q_cur) for comp in nullspace_components]
+                    )
+                    q_cur += alpha * (
+                        J.T @ (np.linalg.solve(jjt, error - J @ (nullspace_term)))
+                        + nullspace_term
+                    )
+
+                n_iters += 1
+
+                if self.visualizer:
+                    self.visualizer.display(q_cur)
+                    time.sleep(VIZ_SLEEP_TIME)
+
+            # Check results at the end of this try
+            if solved:
+                if verbose:
+                    print(f"Solved in {n_tries+1} tries.")
+                break
+            else:
+                q_cur = get_random_state(self.model)
+                n_tries += 1
+                if verbose:
+                    print(f"Retry {n_tries}")
+
+        # Check final results
+        if solved:
+            return q_cur
+        else:
+            return None