scikit-robot 0.4.1__tar.gz → 0.4.2__tar.gz

{scikit_robot-0.4.1/scikit_robot.egg-info → scikit_robot-0.4.2}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scikit-robot
-Version: 0.4.1
+Version: 0.4.2
 Summary: A Flexible Framework for Robot Control in Python
 Author-email: iory <ab.ioryz@gmail.com>
 License: MIT

{scikit_robot-0.4.1 → scikit_robot-0.4.2}/pyproject.toml

@@ -4,7 +4,7 @@ build-backend = "setuptools.build_meta"
 
 [project]
 name = "scikit-robot"
-version = "0.4.1"
+version = "0.4.2"
 description = "A Flexible Framework for Robot Control in Python"
 readme = "README.md"
 license = {text = "MIT"}

{scikit_robot-0.4.1 → scikit_robot-0.4.2/scikit_robot.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: scikit-robot
-Version: 0.4.1
+Version: 0.4.2
 Summary: A Flexible Framework for Robot Control in Python
 Author-email: iory <ab.ioryz@gmail.com>
 License: MIT

{scikit_robot-0.4.1 → scikit_robot-0.4.2}/scikit_robot.egg-info/SOURCES.txt

@@ -110,7 +110,10 @@ skrobot/interfaces/ros2/base.py
 skrobot/interfaces/ros2/gripper.py
 skrobot/interfaces/ros2/panda.py
 skrobot/kinematics/__init__.py
+skrobot/kinematics/_base_placement_jax.py
+skrobot/kinematics/base_placement.py
 skrobot/kinematics/differentiable.py
+skrobot/kinematics/inverse_reachability_map.py
 skrobot/kinematics/loop_closure.py
 skrobot/kinematics/reachability_map.py
 skrobot/model/__init__.py

scikit_robot-0.4.2/skrobot/backend/math_utils.py

@@ -0,0 +1,78 @@
+"""Backend-agnostic math utilities for differentiable computation.
+
+This module provides common mathematical operations that work with
+any backend (NumPy, JAX) for use in kinematics and dynamics computations.
+"""
+
+import numpy as np
+
+
+# _CROSS[i, j, k] = -epsilon_ijk, so _CROSS @ v is the cross-product
+# matrix of v. One matmul builds it for a single vector or a batch; nested
+# stacks would cost several times as much per call on the NumPy FK path.
+_CROSS = np.zeros((3, 3, 3))
+for _i, _j, _k in ((0, 1, 2), (1, 2, 0), (2, 0, 1)):
+    _CROSS[_i, _j, _k] = -1.0
+    _CROSS[_i, _k, _j] = 1.0
+
+
+def rodrigues_rotation(backend, axis, angle):
+    """Compute rotation matrix from axis-angle using Rodrigues' formula.
+
+    R = I + sin(θ)K + (1-cos(θ))K²
+
+    where K is the skew-symmetric matrix of the normalized axis.
+
+    Parameters
+    ----------
+    backend : DifferentiableBackend or module
+        Backend to use for computation (numpy or jax). The ``numpy`` and
+        ``jax.numpy`` modules work as well.
+    axis : array, shape (3,) or (..., 3)
+        Rotation axis (will be normalized), or a batch of them.
+    angle : float or array, shape () or (...)
+        Rotation angle in radians, one per axis.
+
+    Returns
+    -------
+    rotation : array, shape (3, 3) or (..., 3, 3)
+        Rotation matrix, or one per axis.
+
+    Examples
+    --------
+    >>> from skrobot.backend import get_backend
+    >>> from skrobot.backend.math_utils import rodrigues_rotation
+    >>> backend = get_backend('numpy')
+    >>> axis = backend.array([0.0, 0.0, 1.0])
+    >>> R = rodrigues_rotation(backend, axis, 0.5)
+    """
+    # Normalize axis
+    axis = axis / (backend.sqrt((axis ** 2).sum(axis=-1, keepdims=True))
+                   + 1e-10)
+    K = skew_symmetric(backend, axis)
+
+    # Rodrigues formula
+    sin = backend.sin(angle)[..., None, None]
+    cos = backend.cos(angle)[..., None, None]
+    return backend.eye(3) + sin * K + (1.0 - cos) * (K @ K)
+
+
+def skew_symmetric(backend, v):
+    """Create skew-symmetric matrix from a 3D vector.
+
+    The skew-symmetric matrix [v]_× satisfies: [v]_× @ u = v × u
+
+    Parameters
+    ----------
+    backend : DifferentiableBackend or module
+        Backend to use for computation. The ``numpy`` and ``jax.numpy``
+        modules work as well.
+    v : array, shape (3,) or (..., 3)
+        Input vector, or a batch of them.
+
+    Returns
+    -------
+    K : array, shape (3, 3) or (..., 3, 3)
+        Skew-symmetric matrix, or one per vector.
+    """
+    return (_CROSS @ v[..., None, :, None])[..., 0]

{scikit_robot-0.4.1 → scikit_robot-0.4.2}/skrobot/coordinates/base.py

@@ -1794,6 +1794,7 @@ class CascadedCoords(Coordinates):
                     "`relative_coords`'s type should be"
                     "skrobot.coordinates.Coordinates, but is {}"
                     .format(type(relative_coords)))
+            # The parent setter removes child from its previous parent.
             child.parent = self
             child.newcoords(relative_coords, check_validity=False, relative_coords='local')
             self._descendants.append(child)
@@ -1997,7 +1998,8 @@ class CascadedCoords(Coordinates):
         """Rotate this coordinate.
 
         Rotate this coordinate relative to axis by theta radians
-        with respect to wrt.
+        with respect to wrt. The origin of this coordinate is kept for
+        every ``wrt``; only its orientation changes.
 
         Parameters
         ----------
@@ -2006,6 +2008,13 @@ class CascadedCoords(Coordinates):
         axis : str or numpy.ndarray
             'x', 'y', 'z' or vector
         wrt : str or Coordinates
+            Frame in which ``axis`` is expressed.
+
+            - ``'local'`` or ``self``: this coordinate's own axes.
+            - ``'parent'`` or the parent object: the parent's axes. Without
+              a parent, this is the world frame.
+            - ``'world'``: the world axes.
+            - other ``Coordinates``: that coordinate's axes.
         skip_normalization : bool
             if `True`, skip normalization for axis.
 
@@ -2043,22 +2052,9 @@ class CascadedCoords(Coordinates):
         if isinstance(axis, np.ndarray) and axis.shape == (3, 3):
             return self.rotate_with_matrix(theta, wrt)
 
-        if wrt == 'local' or wrt == self:
-            rotation = rotate_matrix_by_axis_angle(
-                self._rotation, theta, axis,
-                skip_normalization=skip_normalization)
-            return self.newcoords(rotation, self._translation,
-                                  check_validity=False, relative_coords='local')
-        elif wrt == 'parent' or wrt == self.parent:
-            rotation = rotate_matrix_by_axis_angle(
-                self._rotation, theta, axis,
-                skip_normalization=skip_normalization)
-            return self.newcoords(rotation, self._translation,
-                                  check_validity=False, relative_coords='local')
-        else:
-            return self.rotate_with_matrix(
-                rotation_matrix(theta, axis,
-                                skip_normalization=skip_normalization), wrt)
+        return self.rotate_with_matrix(
+            rotation_matrix(theta, axis,
+                            skip_normalization=skip_normalization), wrt)
 
     def orient_with_matrix(self, rotation_matrix, wrt='world'):
         """Force update this coordinate system's rotation.
@@ -2192,7 +2188,27 @@ class CascadedCoords(Coordinates):
         if not (c is None or coordinates_p(c)):
             raise ValueError('parent should be None or Coordinates. '
                              'get type=={}'.format(type(c)))
+        previous = getattr(self, '_parent', None)
+        if previous is c:
+            self._parent = c
+            return
+        # A parent chain through self would make worldcoords() recurse
+        # forever.
+        ancestor = c
+        while ancestor is not None:
+            if ancestor is self:
+                raise ValueError(
+                    'Cannot make {!r} a parent of {!r}: it is already one '
+                    'of its descendants.'.format(c.name, self.name))
+            ancestor = getattr(ancestor, 'parent', None)
+        # Leave the previous parent's descendants, so that re-parenting by
+        # either assoc or a direct assignment leaves no stale entry there.
+        if (isinstance(previous, CascadedCoords)
+                and self in previous._descendants):
+            previous._descendants.remove(self)
         self._parent = c
+        # The cached world pose was taken under the previous parent.
+        self.changed()
 
     def __getstate__(self):
         assert self._worldcoords._hook == self.update

{scikit_robot-0.4.1 → scikit_robot-0.4.2}/skrobot/kinematics/__init__.py

@@ -11,6 +11,12 @@ The algorithms support multiple backends (NumPy, JAX) for
 different performance and autodiff requirements.
 """
 
+# Mobile-base placement for several end-effector targets
+from skrobot.kinematics.base_placement import BasePlacementResult
+from skrobot.kinematics.base_placement import find_base_placement
+from skrobot.kinematics.base_placement import set_base_pose
+from skrobot.kinematics.base_placement import vectorized_target
+
 # Differentiable kinematics (backend-agnostic)
 from skrobot.kinematics.differentiable import create_batch_ik_solver
 from skrobot.kinematics.differentiable import create_dynamic_limit_mask
@@ -18,6 +24,10 @@ from skrobot.kinematics.differentiable import extract_fk_parameters
 from skrobot.kinematics.differentiable import forward_kinematics
 from skrobot.kinematics.differentiable import forward_kinematics_ee
 
+# Inverse reachability map: base poses that reach a hand pose
+from skrobot.kinematics.inverse_reachability_map import BaseCandidates
+from skrobot.kinematics.inverse_reachability_map import InverseReachabilityMap
+
 # Loop-closure solving for cut kinematic loops
 from skrobot.kinematics.loop_closure import LoopClosureSolver
 
@@ -33,4 +43,10 @@ __all__ = [
     'create_dynamic_limit_mask',
     'LoopClosureSolver',
     'ReachabilityMap',
+    'BasePlacementResult',
+    'find_base_placement',
+    'set_base_pose',
+    'vectorized_target',
+    'BaseCandidates',
+    'InverseReachabilityMap',
 ]

scikit_robot-0.4.2/skrobot/kinematics/_base_placement_jax.py

@@ -0,0 +1,268 @@
+"""JAX version of the shared-base Levenberg-Marquardt loop.
+
+The NumPy loop in :mod:`skrobot.kinematics.base_placement` drops finished
+starts from the batch every iteration, which keeps the work small but
+spends most of its time in Python and in small ``einsum`` calls. Here one
+iteration -- linearize, Schur step, trial residual, accept or reject --
+is a single jitted function over a fixed-size batch. The batch shrinks to
+the next power of sixteen when fifteen sixteenths of it have finished,
+so a solve compiles for two or three sizes only; each costs seconds,
+which is why the sizes are spaced so far apart.
+
+Everything runs in float64, inside ``jax.enable_x64`` so the global
+setting is left alone. In float32 the accept and stall tests round
+differently, and the solve takes a different path than the NumPy one.
+
+Targets that follow the base are evaluated on the host, as in the NumPy
+loop, together with their finite-difference derivative; within one
+iteration the trial step moves them along that derivative.
+"""
+
+import numpy as np
+
+from skrobot.backend.math_utils import rodrigues_rotation
+
+
+def _enable_x64(jax):
+    """Context manager turning on float64 for the code inside it.
+
+    It is ``jax.enable_x64`` in current JAX and only
+    ``jax.experimental.enable_x64`` up to 0.7.
+    """
+    if hasattr(jax, 'enable_x64'):
+        return jax.enable_x64(True)
+    from jax.experimental import enable_x64
+    return enable_x64(True)
+
+
+class JaxLevenbergMarquardt(object):
+    """Jitted LM iterations for one :class:`_SharedBaseProblem`."""
+
+    def __init__(self, problem, thre, rthre, max_pos_step, max_rot_step,
+                 bounds):
+        import jax
+        import jax.numpy as jnp
+
+        from skrobot.kinematics.base_placement import _clip_joints
+        from skrobot.kinematics.base_placement import _clip_rows
+        from skrobot.kinematics.base_placement import _cost
+        from skrobot.kinematics.base_placement import _margin_terms
+        from skrobot.kinematics.base_placement import _place_on_base
+        from skrobot.kinematics.base_placement import _pose_residual
+        from skrobot.kinematics.base_placement import _residual_jacobians
+        from skrobot.kinematics.base_placement import _schur_step
+        from skrobot.kinematics.differentiable import compute_geometric_jacobian_jax
+
+        self.jax = jax
+        self.jnp = jnp
+        self.problem = problem
+        fk = problem.fk_params
+        n_joints = fk['n_joints']
+        non_mimic = problem.non_mimic
+        n_opt = problem.n_opt
+        n_targets = problem.n_targets
+
+        jac_one = jax.vmap(lambda a: compute_geometric_jacobian_jax(a, fk))
+        row_mask = jnp.asarray(problem.row_mask)
+        tolerance = (None if problem.rotation_tolerance is None
+                     else jnp.asarray(problem.rotation_tolerance))
+        # Kept in NumPy: only its pattern is read, at trace time.
+        base_frame = problem.base_frame
+        lower = jnp.asarray(problem.lower)
+        upper = jnp.asarray(problem.upper)
+        unbounded = jnp.asarray(problem.unbounded)
+        bounded = jnp.asarray(problem.bounded)
+        span = jnp.asarray(problem.span)
+        margin_weight = problem.margin_weight
+        has_moving = bool(problem.target_fns)
+        moving_idx = jnp.asarray([i for i, _ in problem.target_fns],
+                                 dtype=jnp.int32)
+        # The fixed targets stay on the device; only the few that follow
+        # the base cross over each iteration.
+        static_pos = jnp.asarray(problem.target_pos)
+        static_rot = jnp.asarray(problem.target_rot)
+        lo = None if bounds is None else jnp.asarray(bounds[0])
+        hi = None if bounds is None else jnp.asarray(bounds[1])
+
+        def chain(q):
+            n_cand = q.shape[0]
+            angles = jnp.zeros((n_cand * n_targets, n_joints), q.dtype)
+            angles = angles.at[:, non_mimic].set(q.reshape(-1, n_opt))
+            jac, pos, rot = jac_one(angles)
+            return (jac.reshape(n_cand, n_targets, 6, n_opt),
+                    pos.reshape(n_cand, n_targets, 3),
+                    rot.reshape(n_cand, n_targets, 3, 3))
+
+        def residual(meas_pos, meas_rot, target_pos, target_rot):
+            return _pose_residual(meas_pos, meas_rot, target_pos, target_rot, row_mask,
+                                  tolerance, jnp)
+
+        def linearize(bases, q, target_pos, target_rot, dt):
+            jac, pos, rot = chain(q)
+            (meas_pos, meas_rot, world_pos, world_rot, base_rot,
+             base_pos) = _place_on_base(pos, rot, bases, base_frame, jnp)
+            err = residual(meas_pos, meas_rot, target_pos, target_rot)
+            base_jac, arm_jac = _residual_jacobians(
+                err, jac, world_pos, world_rot, base_rot, base_pos,
+                base_frame, row_mask, tolerance, dt, jnp)
+            return err, base_jac, arm_jac
+
+        def margin_terms(q, floor):
+            return _margin_terms(q, floor, lower, upper, span, bounded,
+                                 margin_weight, jnp)
+
+        def iterate(state, base_weight, mpos, mrot, mdt):
+            bases, q, damping, stalls, done, weight, floor = state
+            n_cand = bases.shape[0]
+            target_pos = jnp.broadcast_to(static_pos, (n_cand,) + static_pos.shape)
+            target_rot = jnp.broadcast_to(static_rot, (n_cand,) + static_rot.shape)
+            dt = None
+            if has_moving:
+                target_pos = target_pos.at[:, moving_idx].set(mpos)
+                target_rot = target_rot.at[:, moving_idx].set(mrot)
+                dt = jnp.zeros((n_cand, n_targets, 6, 3)).at[
+                    :, moving_idx].set(mdt)
+            err, bj, aj = linearize(bases, q, target_pos, target_rot, dt)
+            m_err, m_jac = margin_terms(q, floor)
+            cost = _cost(err * weight[..., None], m_err, jnp)
+            pos_norm = jnp.linalg.norm(err[..., :3], axis=-1)
+            rot_norm = jnp.linalg.norm(err[..., 3:], axis=-1)
+            reached = (pos_norm < thre) & (rot_norm < rthre)
+            done = done | (jnp.all(reached | (weight == 0), axis=1)
+                           & ~jnp.any(m_err > 0, axis=(1, 2)))
+            done = done | (damping >= 1e4) | (stalls >= 5)
+            w = weight[..., None]
+            d_base, d_arm = _schur_step(
+                _clip_rows(err, max_pos_step, max_rot_step, jnp) * w,
+                bj * w[..., None], aj * w[..., None], damping, base_weight,
+                m_err, m_jac, jnp)
+            new_bases = bases.at[:, :3].add(d_base)
+            if lo is not None:
+                new_bases = new_bases.at[:, :2].set(
+                    jnp.clip(new_bases[:, :2], lo, hi))
+            new_q = _clip_joints(q + d_arm, lower, upper, unbounded, jnp)
+            step = new_bases[:, :3] - bases[:, :3]
+            if has_moving:
+                # First-order move of the base-following targets.
+                new_target_pos = target_pos + jnp.einsum(
+                    'knia,ka->kni', dt[..., :3, :], step)
+                twist = jnp.einsum('knia,ka->kni', dt[..., 3:, :], step)
+                new_target_rot = rodrigues_rotation(
+                    jnp, twist, jnp.linalg.norm(twist, axis=-1)) @ target_rot
+            else:
+                new_target_pos, new_target_rot = target_pos, target_rot
+            _, pos, rot = chain(new_q)
+            mp, mr, _, _, _, _ = _place_on_base(pos, rot, new_bases,
+                                                base_frame, jnp)
+            new_err = residual(mp, mr, new_target_pos, new_target_rot)
+            new_m_err, _ = margin_terms(new_q, floor)
+            new_cost = _cost(new_err * weight[..., None], new_m_err, jnp)
+            accept = (new_cost < cost) & ~done
+            damping = jnp.where(accept, jnp.maximum(damping * 0.5, 1e-6),
+                                jnp.minimum(damping * 4.0, 1e4))
+            small = cost - new_cost < 1e-3 * jnp.maximum(cost, 1e-12)
+            stalls = jnp.where(accept, jnp.where(small, stalls + 1, 0),
+                               stalls)
+            bases = jnp.where(accept[:, None], new_bases, bases)
+            q = jnp.where(accept[:, None, None], new_q, q)
+            return (bases, q, damping, stalls, done, weight, floor)
+
+        self._iterate = jax.jit(iterate)
+        self._has_moving = has_moving
+
+    def _moving(self, bases):
+        """Host-side poses ``(K, D, 3)``, ``(K, D, 3, 3)`` of the targets
+        that follow the base, and their base derivative ``(K, D, 6, 3)``."""
+        problem = self.problem
+        n_cand = bases.shape[0]
+        if not problem.target_fns:
+            return (np.zeros((n_cand, 0, 3)), np.zeros((n_cand, 0, 3, 3)),
+                    np.zeros((n_cand, 0, 6, 3)))
+        pos, rot = problem._moving_targets(bases)
+        return pos, rot, problem.moving_target_jacobian(bases)
+
+    def run(self, bases, q, active, stop, base_weight, margin_target):
+        """Iterate in place on ``bases`` / ``q``; return the final error."""
+        with _enable_x64(self.jax):
+            self._run(bases, q, active, stop, base_weight, margin_target)
+        return self.problem.residual(bases, q)
+
+    def _run(self, bases, q, active, stop, base_weight, margin_target):
+        n_cand = bases.shape[0]
+        floor = (np.zeros(n_cand) if margin_target is None
+                 else np.asarray(margin_target, dtype=np.float64))
+        weight = active.astype(np.float64)
+        damping = np.full(n_cand, 1e-2)
+        stalls = np.zeros(n_cand, dtype=np.int64)
+        done = np.zeros(n_cand, dtype=bool)
+        base_weight = self.jax.device_put(
+            np.asarray(base_weight, dtype=np.float64))
+
+        jax = self.jax
+        rows = np.arange(n_cand)
+        size = _bucket(n_cand)
+        state = None
+        moving = None
+        for _ in range(stop):
+            if state is None:
+                padded = _pad(rows, size)
+                # One transfer for the lot; jnp.asarray per array costs a
+                # dispatch each.
+                state = jax.device_put((
+                    bases[padded], q[padded], damping[padded],
+                    stalls[padded], done[padded] | (padded < 0),
+                    weight[padded], floor[padded]))
+                host_bases = bases[padded].copy()
+                host_moving = self._moving(host_bases)
+                moving = jax.device_put(host_moving)
+            elif self._has_moving:
+                # Only rows whose base moved need their targets again; a
+                # rejected step leaves the base where it was.
+                new_bases = np.asarray(state[0])
+                changed = np.any(new_bases != host_bases, axis=1)
+                if np.any(changed):
+                    fresh = self._moving(new_bases[changed])
+                    for held, new in zip(host_moving, fresh):
+                        held[changed] = new
+                    host_bases = new_bases.copy()
+                    moving = jax.device_put(host_moving)
+            state = self._iterate(state, base_weight, *moving)
+            live = ~np.asarray(state[4])[:len(rows)]
+            if not np.any(live):
+                break
+            if size > 16 and np.count_nonzero(live) <= size // 16:
+                # Shrink the batch: write the rows back and repack the
+                # ones still running.
+                self._write_back(state, rows, bases, q, damping, stalls,
+                                 done)
+                rows = rows[live]
+                size = _bucket(len(rows))
+                state = None
+        if state is not None:
+            self._write_back(state, rows, bases, q, damping, stalls, done)
+
+    @staticmethod
+    def _write_back(state, rows, bases, q, damping, stalls, done):
+        n = len(rows)
+        bases[rows] = np.asarray(state[0])[:n]
+        q[rows] = np.asarray(state[1])[:n]
+        damping[rows] = np.asarray(state[2])[:n]
+        stalls[rows] = np.asarray(state[3])[:n]
+        done[rows] = np.asarray(state[4])[:n]
+
+
+def _bucket(n):
+    """Smallest power of sixteen, at least 16, holding ``n`` rows."""
+    size = 16
+    while size < n:
+        size *= 16
+    return size
+
+
+def _pad(rows, size):
+    """``rows`` padded to ``size`` with -1, which marks a dummy row.
+
+    Index -1 picks the last row of the batch, a real row whose copy is
+    flagged done and never written back.
+    """
+    return np.concatenate([rows, np.full(size - len(rows), -1)])