diff --git a/README.md b/README.md index f2ad38c..7c6fcc8 100644 --- a/README.md +++ b/README.md @@ -20,6 +20,7 @@ Some unofficial versions in other languages are being developed: * [C++ version](https://github.com/Le0nX/ModernRoboticsCpp) * [Julia version](https://github.com/ferrolho/ModernRoboticsBook.jl) * [Nim version](https://github.com/Nimbotics/ModernRoboticsNim) +* [PyTorch version](https://github.com/Agaggar/PyTorch_MR) in `packages/PyTorch`, installable as `pip install -e packages/PyTorch` and used as `import pytorch_mr as pmr` Some libraries built on ours: * [KinematicsFromDescriptionTool](https://github.com/Interbotix/kinematics_from_description), which calculates the kinematics input parameters from a robot's URDF or robot_description parameter using ROS and Python3. diff --git a/packages/PyTorch/.gitignore b/packages/PyTorch/.gitignore new file mode 100644 index 0000000..4b20b14 --- /dev/null +++ b/packages/PyTorch/.gitignore @@ -0,0 +1,8 @@ +dist/ +build/ +*.egg-info/ +.eggs/ +.pytest_cache/ +__pycache__/ +*.pyc +src/*.egg-info/ diff --git a/packages/PyTorch/LICENSE b/packages/PyTorch/LICENSE new file mode 100644 index 0000000..24780bb --- /dev/null +++ b/packages/PyTorch/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Northwestern University CRB + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/packages/PyTorch/README.md b/packages/PyTorch/README.md new file mode 100644 index 0000000..c47f197 --- /dev/null +++ b/packages/PyTorch/README.md @@ -0,0 +1,58 @@ +# PyTorch Modern Robotics (`pytorch_mr`) + +PyTorch implementation of the routines from *Modern Robotics: Mechanics, Planning, and Control*, aligned with the NumPy `modern_robotics` reference in `packages/Python`. + +## Install (global) +```bash +pip install pytorch-mr +``` + +## Install (local / editable) + +From this directory (`packages/PyTorch`): + +```bash +pip install -e . +``` + +Dependencies are declared in `pyproject.toml` (`numpy`, `torch`). + +## Import + +```python +import pytorch_mr as pmr +# or +from pytorch_mr.core import MatrixExp3, FKinBody +``` + +## Tests (NumPy parity) + +Install the reference `modern_robotics` (local tree or PyPI) and `pytest`: + +```bash +pip install -e ../Python # local reference +# or +pip install modern_robotics +pip install pytest +``` + +From `packages/PyTorch` (uses `[tool.pytest.ini_options]` `pythonpath = ["src"]`): + +```bash +cd packages/PyTorch +python -m pytest tests -q +``` + +From `packages/` with explicit `PYTHONPATH`: + +```bash +cd packages +PYTHONPATH=PyTorch/src:Python python -m pytest PyTorch/tests -q +``` + +Chapter summary: + +```bash +cd packages +PYTHONPATH=PyTorch/src:Python python PyTorch/tests/run_tests_by_chapter.py +``` diff --git a/packages/PyTorch/pyproject.toml b/packages/PyTorch/pyproject.toml new file mode 100644 index 0000000..8ef6e38 --- /dev/null +++ b/packages/PyTorch/pyproject.toml @@ -0,0 +1,51 @@ +[build-system] +requires = ["setuptools>=61", "wheel"] +build-backend = "setuptools.build_meta" + +[project] +name = "pytorch_mr" +description = "PyTorch implementation of Modern Robotics: Mechanics, Planning, and Control" +readme = "README.md" +requires-python = ">=3.6" +license = { file = "LICENSE" } +authors = [ + { name = "Ayush Gaggar", email = "agaggar@u.northwestern.edu" } +] +keywords = ["robotics", "kinematics", "dynamics", "pytorch", "modern robotics"] +classifiers = [ + "Development Status :: 3 - Alpha", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Topic :: Education", + "Topic :: Scientific/Engineering", + "Programming Language :: Python :: 3", + "Operating System :: OS Independent", +] +dependencies = [ + "numpy>=1.20", + "torch>=1.10", +] +dynamic = ["version"] + +[project.urls] +Homepage = "https://github.com/Agaggar/PyTorch_MR/tree/master/packages/PyTorch" +Documentation = "http://modernrobotics.org/" +Repository = "https://github.com/Agaggar/PyTorch_MR.git" +Issues = "https://github.com/Agaggar/PyTorch_MR/issues" + +[project.optional-dependencies] +dev = ["pytest>=7"] + +[tool.setuptools.packages.find] +where = ["src"] + +# Read version from the lightweight module only (avoid importing core.py -> torch in isolated builds). +[tool.setuptools.dynamic] +version = { attr = "pytorch_mr.__version__.__version__" } + +[tool.pytest.ini_options] +testpaths = ["tests"] +pythonpath = ["src"] diff --git a/packages/PyTorch/setup.py b/packages/PyTorch/setup.py new file mode 100644 index 0000000..867027e --- /dev/null +++ b/packages/PyTorch/setup.py @@ -0,0 +1,46 @@ +from setuptools import find_packages, setup + +exec(open('src/pytorch_mr/__version__.py').read()) + +long_description = """ +# Modern Robotics (PyTorch) + +PyTorch implementation of the mathematical routines from +_Modern Robotics: Mechanics, Planning, and Control_. + +This package mirrors the reference NumPy `modern_robotics` library with batched tensor support. + +Install the NumPy reference separately from `packages/Python` if you need parity tests or side-by-side use: +`pip install -e ../Python` (from this directory's parent) or `pip install modern_robotics`. +""" + +setup( + name="pytorch_mr", + version=__version__, + author="Ayush Gaggar, Huan Weng, Mikhail Todes, Jarvis Schultz, Bill Hunt, ", + author_email="agaggar@u.northwestern.edu", + description=( + "Modern Robotics: Mechanics, Planning, and Control — PyTorch implementation" + ), + license="MIT", + long_description=long_description, + long_description_content_type="text/markdown", + keywords="kinematics robotics dynamics pytorch", + url="http://modernrobotics.org/", + packages=find_packages(include=["pytorch_mr", "pytorch_mr.*"]), + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Programming Language :: Python :: 3", + "Topic :: Education", + "Topic :: Scientific/Engineering", + ], + install_requires=[ + "numpy", + "torch", + ], + platforms="Linux, Mac, Windows", +) \ No newline at end of file diff --git a/packages/PyTorch/src/pytorch_mr/__init__.py b/packages/PyTorch/src/pytorch_mr/__init__.py new file mode 100644 index 0000000..aada3de --- /dev/null +++ b/packages/PyTorch/src/pytorch_mr/__init__.py @@ -0,0 +1,3 @@ +from .__version__ import __version__ + +from .core import * diff --git a/packages/PyTorch/src/pytorch_mr/__version__.py b/packages/PyTorch/src/pytorch_mr/__version__.py new file mode 100644 index 0000000..2192205 --- /dev/null +++ b/packages/PyTorch/src/pytorch_mr/__version__.py @@ -0,0 +1,2 @@ + +__version__ = '0.0.1' diff --git a/packages/PyTorch/src/pytorch_mr/core.py b/packages/PyTorch/src/pytorch_mr/core.py new file mode 100644 index 0000000..6025f50 --- /dev/null +++ b/packages/PyTorch/src/pytorch_mr/core.py @@ -0,0 +1,1023 @@ +import torch +''' +*************************************************************************** +Modern Robotics: Mechanics, Planning, and Control. +Code Library for PyTorch +*************************************************************************** +Author: Ayush Gaggar + Adapted from original Python code by Huan Weng, Bill Hunt, Jarvis Schultz, Mikhail Todes, Kevin Lynch, and Frank Park +Email: agaggar@u.northwestern.edu +Date: July 2025 +*************************************************************************** +Language: Python +Required library: pytorch +*************************************************************************** + +This code is a PyTorch implementation of the mathematical operations and transformations in the Modern Robotics package, +and so is best used for *batched* tensors, i.e., tensors with shape (N, ...), where N is the batch size. +For functions after chapter 3, many of the docstrings were written by AI. Code review and shape broadcasting was also completed by AI. +''' + +def _as_batch_vec(x: torch.Tensor, dim: int): + """Ensures vector-like inputs have a batch dimension. + + :param x: Tensor with shape (dim,), (N, dim), or (N, dim, 1) + :param dim: Expected vector dimension + :return: Tensor with shape (N, dim) + """ + if x.ndim == 1: + return x.unsqueeze(0) + if x.ndim == 3 and x.shape[-1] == 1 and x.shape[-2] == dim: + return x.squeeze(-1) + return x + +def _as_batch_mat(x: torch.Tensor, rows: int, cols: int): + """Ensures matrix-like inputs have a batch dimension. + + :param x: Tensor with shape (rows, cols) or (N, rows, cols) + :param rows: Number of rows of the matrix + :param cols: Number of columns of the matrix + :return: Tensor with shape (N, rows, cols) + """ + if x.ndim == 2 and x.shape == (rows, cols): + return x.unsqueeze(0) + return x + +def NearZero(val: torch.Tensor, eps=1e-6): + """Check if input tensor values are close to zero.""" + return torch.abs(val) < eps + +def Normalize(v: torch.Tensor): + """Normalizes a vector to unit length. + + :param v: A vector, shape (N, 3, 1), (N, 3), (3,) + :return: A normalized vector, shape (N, 3) + """ + if v.ndim == 1: + v = v.unsqueeze(0) + if v.ndim == 2: + v = v.unsqueeze(-1) + norm = torch.linalg.norm(v, dim=-2, keepdim=True) + return v / norm.clamp(min=1e-6) # Avoid division by zero + +def RotInv(R: torch.Tensor): + """Inverts a rotation matrix. + + :param R: A rotation matrix, shape (N, 3, 3) + :return: The inverse of R, which is the transpose of R, shape (N, 3, 3) + """ + if R.ndim == 2: + R = R.unsqueeze(0) + return torch.transpose(R, -2, -1) + +def VecToso3(omega: torch.Tensor): + """Converts a 3-vector to an so(3) representation + + :param omg: angular velocity vector, shape (N, 3) + :return: The skew symmetric representation of omg, i.e., an so(3) matrix, + shape (N, 3, 3) + """ + if omega.ndim > 2: + omega = omega.reshape(-1, 3) + elif omega.ndim == 1: + omega = omega.unsqueeze(0) + wx, wy, wz = omega[:, 0], omega[:, 1], omega[:, 2] + row0 = torch.stack([torch.zeros_like(wx), -wz, wy], dim=1) + row1 = torch.stack([wz, torch.zeros_like(wy), -wx], dim=1) + row2 = torch.stack([-wy, wx, torch.zeros_like(wz)], dim=1) + so3 = torch.stack([row0, row1, row2], dim=1) + return so3 + +def so3ToVec(so3mat: torch.Tensor): + """Converts an so(3) representation to a 3-vector + + :param so3mat: A skew-symmetric matrix, shape (N, 3, 3) + :return: The 3-vector corresponding to so3mat, shape (N, 3, 1) + """ + if so3mat.ndim == 2: + so3mat = so3mat.unsqueeze(0) + return torch.stack([ + so3mat[:, 2, 1], + so3mat[:, 0, 2], + so3mat[:, 1, 0] + ], dim=-1).unsqueeze(-1) + + +def AxisAng3(expc3: torch.Tensor): + """Converts a 3-vector of exponential coordinates for rotation into + axis-angle form + + :param expc3: A 3-vector of exponential coordinates for rotation, shape (N, 3, 1) + :return omghat: A unit rotation axis, shape (N, 3) + :return theta: The corresponding rotation angle, shape (N,) + """ + expc3 = _as_batch_vec(expc3, 3) + if expc3.ndim == 2: + expc3 = expc3.unsqueeze(-1) # (N,3,1) + theta = torch.linalg.norm(expc3, dim=-2).reshape(expc3.shape[0]) + omega_hat = torch.zeros_like(expc3) + nonzero = ~NearZero(theta) + omega_hat[nonzero] = expc3[nonzero] / theta[nonzero][:, None, None] + return omega_hat, theta + +def MatrixExp3(so3mat: torch.Tensor): + """Computes the matrix exponential of a matrix in so(3) + + :param so3mat: A skew-symmetric matrix, shape (N, 3, 3) + :return: The matrix exponential of so3mat, shape (N, 3, 3) + + Example Input: + so3mat = np.array([[ 0, -3, 2], + [ 3, 0, -1], + [-2, 1, 0]]) + Output: + np.array([[-0.69492056, 0.71352099, 0.08929286], + [-0.19200697, -0.30378504, 0.93319235], + [ 0.69297817, 0.6313497 , 0.34810748]]) + """ + so3mat = _as_batch_mat(so3mat, 3, 3) + omgtheta = so3ToVec(so3mat) + nonzero = ~NearZero(torch.linalg.norm(omgtheta, dim=-2).reshape(so3mat.shape[0])) + add_me = torch.zeros_like(so3mat) + theta_nonzero = AxisAng3(omgtheta)[1][nonzero][:, None, None] + omgmat_nonzero = so3mat[nonzero] / theta_nonzero + add_me[nonzero] = torch.sin(theta_nonzero) * omgmat_nonzero + (1 - torch.cos(theta_nonzero)) * torch.bmm(omgmat_nonzero, omgmat_nonzero) + identity = torch.eye(3, device=so3mat.device, dtype=so3mat.dtype) + identity = identity.repeat(so3mat.shape[0], 1, 1) + return identity + add_me + +def MatrixLog3(R: torch.Tensor): + """Computes the matrix logarithm of rotation matrix + + :param R: Nx3x3 rotation matrix, shape (N, 3, 3) + :return: The matrix logarithm of R, shape (N, 3, 3) + """ + R = _as_batch_mat(R, 3, 3) + #TODO: broadcast instead of iterating + trace_R = R.diagonal(offset=0, dim1=1, dim2=2).sum(-1) # (N,) + acosinput = torch.clamp((trace_R - 1) / 2.0, -1.0, 1.0) + logm = torch.zeros_like(R, device=R.device, dtype=R.dtype) + for idx, rot in enumerate(R): + if torch.abs(1 - acosinput[idx]) <= 1e-6: + continue + elif torch.abs(acosinput[idx] + 1) <= 1e-6: + if not NearZero(1 + rot[2][2]): + omg = (1.0 / torch.sqrt(torch.clip(2 * (1 + rot[2][2]), min=1e-6))) \ + * torch.stack([rot[0][2], rot[1][2], 1 + rot[2][2]], dim=0) + elif not NearZero(1 + rot[1][1]): + omg = (1.0 / torch.sqrt(torch.clip(2 * (1 + rot[1][1]), min=1e-6))) \ + * torch.stack([rot[0][1], 1 + rot[1][1], rot[2][1]], dim=0) + else: + omg = (1.0 / torch.sqrt(torch.clip(2 * (1 + rot[0][0]), min=1e-6))) \ + * torch.stack([1 + rot[0][0], rot[1][0], rot[2][0]], dim=0) + logm[idx] = VecToso3(torch.pi * omg) + else: + theta = torch.acos(acosinput[idx]) + logm[idx] = theta / (2.0 * torch.sin(theta)) * (rot - torch.transpose(rot, -2, -1)) + if torch.any(torch.isnan(logm)): + raise ValueError("NaN values in MatrixLog3") + return logm + +def RpToTrans(R: torch.Tensor, p: torch.Tensor): + """Converts rotation matrices and position vectors into homogeneous + transformation matrices + + :param R: Rotation matrix, (N, 3, 3) + :param p: A 3-vector, (N, 3) + :return: A homogeneous transformation matrix corresponding to the inputs + """ + if p.ndim == 2: + p = p.unsqueeze(-1) + bottom_row = torch.tensor([0, 0, 0, 1], device=R.device, dtype=R.dtype) + bottom_row = bottom_row.expand(R.size(0), 1, 4) + return torch.cat([torch.cat([R, p], dim=-1), + bottom_row], dim=-2).float().to(device=R.device) + +def TransToRp(T: torch.Tensor): + """Converts a homogeneous transformation matrix into a rotation matrix + and position vector, given T is (N, 4, 4). + + :param T: A homogeneous transformation matrix + :return R: The corresponding rotation matrix, (N, 3, 3) + :return p: The corresponding position vector, (N, 3, 1) + """ + if T.ndim == 2: + T = T.unsqueeze(0) + return T[:, 0: 3, 0: 3], T[:, 0: 3, 3].unsqueeze(-1) + +def TransInv(T: torch.Tensor): + """Inverts a homogeneous transformation matrix, given T is (N, 4, 4). + + :param T: A homogeneous transformation matrix, (N, 4, 4) + :return: The inverse of T, (N, 4, 4) + Uses the structure of transformation matrices to avoid taking a matrix + inverse, for efficiency. + """ + R, p = TransToRp(T) + Rt = torch.transpose(R, -2, -1) + if p.ndim == 2: + p = p.unsqueeze(-1) + bottom_row = torch.tensor([0, 0, 0, 1], device=R.device, dtype=R.dtype) + bottom_row = bottom_row.expand(R.size(0), 1, 4) + return torch.cat([torch.cat([Rt, -torch.bmm(Rt, p)], dim=-1), + bottom_row], dim=-2) + +def VecTose3(V: torch.Tensor): + """Converts a spatial velocity vector into a 4x4 matrix in se3 + + :param V: A 6-vector representing a spatial velocity, shape (N, 6) + :return: The 4x4 se3 representation of V, shape (N, 4, 4) + """ + omega = V[:, :3] + v = V[:, 3:] + if v.ndim == 2: + v = v.unsqueeze(-1) + omega = omega.unsqueeze(-1) + # bottom_row = torch.zeros((V.shape[0], 1, 4), device=V.device, dtype=V.dtype) + # bottom_row[:, 0, 3] = 1. + bottom_row = torch.zeros((V.size(0), 1, 4), device=V.device, dtype=V.dtype) + return torch.cat([torch.cat([VecToso3(omega), v], dim=-1), + bottom_row], dim=-2) + +def se3ToVec(se3mat: torch.Tensor): + """ Converts se3 matrices into spatial velocity vector + + :param se3mat: SE3 Matrix, shape (N, 4, 4) + :return: The spatial velocity 6-vector corresponding to se3mat, shape (N, 6, 1) + """ + if se3mat.ndim == 2: + se3mat = se3mat.unsqueeze(0) + vecm = torch.stack([ + se3mat[:, 2, 1], + se3mat[:, 0, 2], + se3mat[:, 1, 0], + se3mat[:, 0, 3], + se3mat[:, 1, 3], + se3mat[:, 2, 3], + ], dim=1) + return vecm.unsqueeze(-1) + +def Adjoint(T: torch.Tensor): + """Computes the adjoint representation of a homogeneous transformation + matrix + + :param T: A homogeneous transformation matrix, shape (N, 4, 4) + :return: The adjoint representation of T, shape (N, 6, 6) + """ + R, p = TransToRp(T) + top_left = R + top_right = torch.zeros((R.shape[0], 3, 3), device=T.device, dtype=T.dtype) + bottom_left = torch.bmm(VecToso3(p), R) + bottom_right = R + return torch.cat([torch.cat([top_left, top_right], dim=-1), + torch.cat([bottom_left, bottom_right], dim=-1)], dim=-2) + +def ScrewToAxis(q: torch.Tensor, s: torch.Tensor, h: torch.Tensor): + """Takes a parametric description of a screw axis and converts it to a + normalized screw axis. + + :param q: A point on the screw axis, shape (N, 3) + :param s: The direction of the screw axis, shape (N, 3) + :param h: The pitch of the screw, shape (N,) + :return: The axis representation of the screw, shape (N, 6) + """ + q = _as_batch_vec(q, 3) + s = _as_batch_vec(s, 3) + if h.ndim == 0: + h = h.unsqueeze(0) + if h.ndim == 2 and h.shape[-1] == 1: + h = h.squeeze(-1) + if h.ndim == 1 and h.shape[0] == 1 and q.shape[0] > 1: + h = h.expand(q.shape[0]) + return torch.cat([s, torch.cross(q, s, dim=-1) + h.unsqueeze(-1) * s], dim=-1) + +def AxisAng6(expc6: torch.Tensor): + """Converts a 6-vector of exponential coordinates into screw axis-angle + form + + :param expc6: A 6-vector of exponential coordinates for rigid-body motion + S*theta, shape (N, 6) or (N, 6, 1) + :return S: The corresponding normalized screw axis, shape (N, 6) + :return theta: The distance traveled along/about S, shape (N, 1) or (N,) + """ + expc6 = _as_batch_vec(expc6, 6) + theta = torch.linalg.norm(expc6[:, :3], dim=-1) + mask = NearZero(theta) + theta[mask] = torch.linalg.norm(expc6[mask, 3:], dim=-1) + theta_safe = torch.where(NearZero(theta), torch.ones_like(theta), theta) + return expc6 / theta_safe.unsqueeze(-1), theta + +def MatrixExp6(se3mat: torch.Tensor): + """Computes the matrix exponential of an se3 representation of + exponential coordinates + + :param se3mat: Matrix in se3, shape (N, 4, 4) + :return: Matrix exponential of se3mat, shape (N, 4, 4) + """ + omgtheta = so3ToVec(se3mat[:, :3, :3]) + nonzero = ~NearZero(torch.linalg.norm(omgtheta, dim=-2).reshape(se3mat.shape[0])) + + first_part = MatrixExp3(se3mat[:, :3, :3]) + second_part = se3mat[:, :3, 3].unsqueeze(-1).clone() + identity = torch.eye(3, device=se3mat.device, dtype=se3mat.dtype) + identity = identity.repeat(se3mat.shape[0], 1, 1) + bottom_row = torch.tensor([0, 0, 0, 1], device=se3mat.device, dtype=se3mat.dtype) + bottom_row = bottom_row.expand(se3mat.size(0), 1, 4) # (N, 1, 4) + + theta_nonzero = AxisAng3(omgtheta)[1][nonzero][:, None, None] + omgmat_nonzero = se3mat[:, :3, :3][nonzero] / theta_nonzero + second_part[nonzero] = torch.bmm(identity[nonzero] * theta_nonzero \ + + (1 - torch.cos(theta_nonzero)) * omgmat_nonzero \ + + (theta_nonzero - torch.sin(theta_nonzero)) \ + * torch.bmm(omgmat_nonzero, omgmat_nonzero), + second_part[nonzero] / theta_nonzero) + omgmat_nonzero = se3mat[:, :3, :3][nonzero] / theta_nonzero + + return torch.cat([torch.cat([first_part, second_part], dim=-1), + bottom_row], dim=-2) + +def MatrixLog6(T: torch.Tensor): + """Computes the matrix logarithm of a homogeneous transformation matrix + + :param T: Matrix in SE3, shape (N, 4, 4) + :return: Matrix logarithm of T, shape (N, 4, 4) + """ + T = _as_batch_mat(T, 4, 4) + rot, second = TransToRp(T) + all_rots = MatrixLog3(rot) + if second.ndim == 2: + second = second.unsqueeze(-1) + + first = torch.zeros_like(rot, device=T.device, dtype=T.dtype) + bottom_row = torch.zeros((T.shape[0], 1, 4), device=T.device, dtype=T.dtype) + logm = torch.zeros_like(T, device=T.device, dtype=T.dtype) + + #TODO: broadcast instead of iterating + nonzero = ~(NearZero(all_rots, eps=1e-4).sum(dim=2).sum(dim=1) == 9) + logm[~nonzero] = torch.cat([torch.cat([first[~nonzero], second[~nonzero]], dim=-1), + bottom_row[~nonzero]], dim=-2) + trace_R = rot.diagonal(offset=0, dim1=1, dim2=2).sum(-1) # (N,) + if torch.any(torch.isnan(trace_R)): + raise ValueError("NaN values in trace_R") + theta_nonzero = torch.acos(torch.clip((trace_R - 1) / 2.0, -1., 1.))[nonzero] + + identity = torch.eye(3, device=T.device, dtype=T.dtype) + identity = identity.repeat(T.shape[0], 1, 1) + logm[nonzero] = torch.cat([torch.cat([all_rots[nonzero], + (identity[nonzero] - all_rots[nonzero] / 2.0 \ + + torch.clamp((1.0 / theta_nonzero - 1.0 / torch.tan(theta_nonzero / 2.0) / 2.), min=-1e3, max=1e4)[:, None, None] \ + * torch.bmm(all_rots[nonzero], all_rots[nonzero]) / theta_nonzero[:, None, None]) @ second[nonzero]], dim=-1), + bottom_row[nonzero]], dim=-2) + if torch.any(torch.isnan(logm)): + raise ValueError("NaN values in MatrixLog6") + return logm + +def ProjectToSO3(R: torch.Tensor): + """Projects a rotation matrix onto the nearest SO(3) matrix + + :param R: A rotation matrix, shape (N, 3, 3) + :return: The projected rotation matrix, shape (N, 3, 3) + """ + R = _as_batch_mat(R, 3, 3) + U, _, Vh = torch.linalg.svd(R) + Rproj = torch.matmul(U, Vh) + det = torch.linalg.det(Rproj) + fix_mask = det < 0 + if torch.any(fix_mask): + Rproj_fix = Rproj[fix_mask].clone() + Rproj_fix[:, :, 2] = -Rproj_fix[:, :, 2] + Rproj[fix_mask] = Rproj_fix + return Rproj + +def ProjectToSE3(T: torch.Tensor): + """Returns a projection of mat into SE(3) + + :param mat: A 4x4 matrix to project to SE(3) + :return: The closest matrix to T that is in SE(3) + Projects a matrix mat to the closest matrix in SE(3) using singular-value + decomposition. + """ + T = _as_batch_mat(T, 4, 4) + out = T.clone() + out[:, :3, :3] = ProjectToSO3(T[:, :3, :3]) + out[:, 3, :] = torch.tensor([0, 0, 0, 1], device=T.device, dtype=T.dtype).expand(T.shape[0], 4) + return out + +def DistanceToSO3(R: torch.Tensor): + """Returns the Frobenius norm to describe the distance of mat from the + SO(3) manifold + + :param R: A matrix, shape (N, 3, 3) + :return: A quantity describing the distance of mat from the SO(3) manifold, shape (N,) + """ + R = _as_batch_mat(R, 3, 3) + det = torch.linalg.det(R) + out = torch.full((R.shape[0],), 1e9, device=R.device, dtype=R.dtype) + mask = det > 0 + if torch.any(mask): + RtR = torch.matmul(torch.transpose(R[mask], -2, -1), R[mask]) + I = torch.eye(3, device=R.device, dtype=R.dtype).unsqueeze(0).expand(RtR.shape[0], 3, 3) + out[mask] = torch.linalg.norm(RtR - I, dim=(1, 2)) + return out + +def DistanceToSE3(T: torch.Tensor): + """Returns the Frobenius norm to describe the distance of mat from the + SE(3) manifold + + :param T: A matrix, shape (N, 4, 4) + :return: A quantity describing the distance of mat from the SE(3) manifold, shape (N,) + + Computes the distance from mat to the SE(3) manifold using the following + method: + Compute the determinant of matR, the top 3x3 submatrix of mat. + If det(matR) <= 0, return a large number. + If det(matR) > 0, replace the top 3x3 submatrix of mat with matR^T.matR, + and set the first three entries of the fourth column of mat to zero. Then + return norm(mat - I). + + Example Input: + T = torch.tensor([[ 1.0, 0.0, 0.0, 1.2 ], + [ 0.0, 0.1, -0.95, 1.5 ], + [ 0.0, 1.0, 0.1, -0.9 ], + [ 0.0, 0.0, 0.1, 0.98 ]], dtype=torch.float64) + Output: + 0.134931 + """ + T = _as_batch_mat(T, 4, 4) + matR = T[:, :3, :3] + det = torch.linalg.det(matR) + out = torch.full((T.shape[0],), 1e9, device=T.device, dtype=T.dtype) + mask = det > 0 + if torch.any(mask): + RtR = torch.matmul(torch.transpose(matR[mask], -2, -1), matR[mask]) + top = torch.cat([RtR, torch.zeros((RtR.shape[0], 3, 1), device=T.device, dtype=T.dtype)], dim=-1) + bottom = T[mask, 3:4, :] + mat = torch.cat([top, bottom], dim=-2) + I = torch.eye(4, device=T.device, dtype=T.dtype).unsqueeze(0).expand(mat.shape[0], 4, 4) + out[mask] = torch.linalg.norm(mat - I, dim=(1, 2)) + return out + +def TestIfSO3(R: torch.Tensor): + """Returns true if mat is close to or on the manifold SO(3) + + :param mat: A 3x3 matrix + :return: True if mat is very close to or in SO(3), false otherwise + """ + return torch.abs(DistanceToSO3(R)) < 1e-3 + +def TestIfSE3(T: torch.Tensor): + """Returns true if mat is close to or on the manifold SE(3) + + :param mat: A 4x4 matrix + :return: True if mat is very close to or in SE(3), false otherwise + """ + return torch.abs(DistanceToSE3(T)) < 1e-3 + +def FKinBody(M: torch.Tensor, Blist: torch.Tensor, thetalist: torch.Tensor): + """Computes forward kinematics in the body frame for an open chain robot + + :param M: The home configuration (position and orientation) of the end-effector + :param Blist: The joint screw axes in the end-effector frame when the + manipulator is at the home position, with axes as columns + :param thetalist: A list of joint coordinates + :return: A homogeneous transformation matrix representing the end-effector + frame when the joints are at the specified coordinates (Body frame) + """ + M = _as_batch_mat(M, 4, 4) + Blist = Blist if Blist.ndim == 3 else Blist.unsqueeze(0) + thetalist = _as_batch_vec(thetalist, Blist.shape[-1]) + T = M.clone() + n = thetalist.shape[-1] + for i in range(n): + T = torch.bmm(T, MatrixExp6(VecTose3(Blist[:, :, i] * thetalist[:, i:i+1]))) + return T + +def FKinSpace(M: torch.Tensor, Slist: torch.Tensor, thetalist: torch.Tensor): + """Computes forward kinematics in the space frame for an open chain robot + + :param M: The home configuration (position and orientation) of the end-effector + :param Slist: The joint screw axes in the space frame when the manipulator + is at the home position, with axes as columns + :param thetalist: A list of joint coordinates + :return: A homogeneous transformation matrix representing the end-effector + frame when the joints are at the specified coordinates (Space frame) + """ + M = _as_batch_mat(M, 4, 4) + Slist = Slist if Slist.ndim == 3 else Slist.unsqueeze(0) + thetalist = _as_batch_vec(thetalist, Slist.shape[-1]) + T = M.clone() + n = thetalist.shape[-1] + for i in range(n - 1, -1, -1): + T = torch.bmm(MatrixExp6(VecTose3(Slist[:, :, i] * thetalist[:, i:i+1])), T) + return T + +def JacobianBody(Blist: torch.Tensor, thetalist: torch.Tensor): + """Computes the body Jacobian for an open chain robot + + :param Blist: The joint screw axes in the end-effector frame when the + manipulator is at the home position, with axes as columns + :param thetalist: A list of joint coordinates + :return: The body Jacobian corresponding to the inputs + """ + Blist = Blist if Blist.ndim == 3 else Blist.unsqueeze(0) + thetalist = _as_batch_vec(thetalist, Blist.shape[-1]) + Jb = Blist.clone().to(dtype=torch.float32 if Blist.dtype in (torch.int32, torch.int64) else Blist.dtype) + T = torch.eye(4, device=Blist.device, dtype=Jb.dtype).unsqueeze(0).repeat(Blist.shape[0], 1, 1) + for i in range(thetalist.shape[-1] - 2, -1, -1): + T = torch.bmm(T, MatrixExp6(VecTose3(Blist[:, :, i + 1] * (-thetalist[:, i + 1:i + 2])))) + Jb[:, :, i] = torch.bmm(Adjoint(T), Blist[:, :, i].unsqueeze(-1)).squeeze(-1) + return Jb + +def JacobianSpace(Slist: torch.Tensor, thetalist: torch.Tensor): + """Computes the space Jacobian for an open chain robot + + :param Slist: The joint screw axes in the space frame when the manipulator + is at the home position, with axes as columns + :param thetalist: A list of joint coordinates + :return: The space Jacobian corresponding to the inputs + """ + Slist = Slist if Slist.ndim == 3 else Slist.unsqueeze(0) + thetalist = _as_batch_vec(thetalist, Slist.shape[-1]) + Js = Slist.clone().to(dtype=torch.float32 if Slist.dtype in (torch.int32, torch.int64) else Slist.dtype) + T = torch.eye(4, device=Slist.device, dtype=Js.dtype).unsqueeze(0).repeat(Slist.shape[0], 1, 1) + for i in range(1, thetalist.shape[-1]): + T = torch.bmm(T, MatrixExp6(VecTose3(Slist[:, :, i - 1] * thetalist[:, i - 1:i]))) + Js[:, :, i] = torch.bmm(Adjoint(T), Slist[:, :, i].unsqueeze(-1)).squeeze(-1) + return Js + +def IKinBody(Blist: torch.Tensor, M: torch.Tensor, T: torch.Tensor, thetalist0: torch.Tensor, eomg=1e-6, ev=1e-6): + """Computes inverse kinematics in the body frame for an open chain robot + + :param Blist: The joint screw axes in the end-effector frame when the + manipulator is at the home position, with axes as columns + :param M: The home configuration of the end-effector + :param T: The desired end-effector configuration Tsd + :param thetalist0: An initial guess of joint angles + :param eomg: Tolerance on orientation error + :param ev: Tolerance on linear position error + :return thetalist: Joint angles that achieve T within the specified tolerances + :return success: True if a solution is found, False otherwise + Uses an iterative Newton-Raphson root-finding method with maxiterations=20. + """ + Blist = Blist if Blist.ndim == 3 else Blist.unsqueeze(0) + M = _as_batch_mat(M, 4, 4) + T = _as_batch_mat(T, 4, 4) + thetalist = _as_batch_vec(thetalist0, Blist.shape[-1]).clone() + i = 0 + maxiterations = 20 + Vb = se3ToVec(MatrixLog6(torch.bmm(TransInv(FKinBody(M, Blist, thetalist)), T))).squeeze(-1) + err = (torch.linalg.norm(Vb[:, :3], dim=-1) > eomg) | (torch.linalg.norm(Vb[:, 3:], dim=-1) > ev) + while torch.any(err) and i < maxiterations: + Jb = JacobianBody(Blist, thetalist) + for b in range(thetalist.shape[0]): + if err[b]: + thetalist[b] = thetalist[b] + torch.matmul(torch.linalg.pinv(Jb[b]), Vb[b]) + i = i + 1 + Vb = se3ToVec(MatrixLog6(torch.bmm(TransInv(FKinBody(M, Blist, thetalist)), T))).squeeze(-1) + err = (torch.linalg.norm(Vb[:, :3], dim=-1) > eomg) | (torch.linalg.norm(Vb[:, 3:], dim=-1) > ev) + return thetalist, ~err + +def IKinSpace(Slist: torch.Tensor, M: torch.Tensor, T: torch.Tensor, thetalist0: torch.Tensor, eomg=1e-6, ev=1e-6): + """Computes inverse kinematics in the space frame for an open chain robot + + :param Slist: The joint screw axes in the space frame when the + manipulator is at the home position, with axes as columns + :param M: The home configuration of the end-effector + :param T: The desired end-effector configuration Tsd + :param thetalist0: An initial guess of joint angles + :param eomg: Tolerance on orientation error + :param ev: Tolerance on linear position error + :return thetalist: Joint angles that achieve T within the specified tolerances + :return success: True if a solution is found, False otherwise + Uses an iterative Newton-Raphson root-finding method with maxiterations=20. + """ + Slist = Slist if Slist.ndim == 3 else Slist.unsqueeze(0) + M = _as_batch_mat(M, 4, 4) + T = _as_batch_mat(T, 4, 4) + thetalist = _as_batch_vec(thetalist0, Slist.shape[-1]).clone() + i = 0 + maxiterations = 20 + Tsb = FKinSpace(M, Slist, thetalist) + Vs = torch.bmm(Adjoint(Tsb), se3ToVec(MatrixLog6(torch.bmm(TransInv(Tsb), T)))).squeeze(-1) + err = (torch.linalg.norm(Vs[:, :3], dim=-1) > eomg) | (torch.linalg.norm(Vs[:, 3:], dim=-1) > ev) + while torch.any(err) and i < maxiterations: + Js = JacobianSpace(Slist, thetalist) + for b in range(thetalist.shape[0]): + if err[b]: + thetalist[b] = thetalist[b] + torch.matmul(torch.linalg.pinv(Js[b]), Vs[b]) + i = i + 1 + Tsb = FKinSpace(M, Slist, thetalist) + Vs = torch.bmm(Adjoint(Tsb), se3ToVec(MatrixLog6(torch.bmm(TransInv(Tsb), T)))).squeeze(-1) + err = (torch.linalg.norm(Vs[:, :3], dim=-1) > eomg) | (torch.linalg.norm(Vs[:, 3:], dim=-1) > ev) + return thetalist, ~err + +def ad(V: torch.Tensor): + """Computes the adjoint representation of a spatial velocity vector + + :param V: A spatial velocity vector, shape (N, 6) + :return: The adjoint representation of V, shape (N, 6, 6) + """ + V = _as_batch_vec(V, 6) + omgmat = VecToso3(V[:, :3]) + zeros = torch.zeros((V.shape[0], 3, 3), device=V.device, dtype=V.dtype) + return torch.cat([torch.cat([omgmat, zeros], dim=-1), + torch.cat([VecToso3(V[:, 3:]), omgmat], dim=-1)], dim=-2) + +def InverseDynamics(thetalist: torch.Tensor, dthetalist: torch.Tensor, ddthetalist: torch.Tensor, g: torch.Tensor, Ftip: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes inverse dynamics in the space frame for an open chain robot + + :param thetalist: n-vector of joint variables + :param dthetalist: n-vector of joint rates + :param ddthetalist: n-vector of joint accelerations + :param g: Gravity vector g + :param Ftip: Spatial force applied by the end-effector expressed in frame {n+1} + :param Mlist: List of link frames {i} relative to {i-1} at the home position + :param Glist: Spatial inertia matrices Gi of the links + :param Slist: Screw axes Si of the joints in a space frame, with axes as columns + :return: The n-vector of required joint forces/torques + This function uses forward-backward Newton-Euler iterations to solve: + taulist = M(thetalist)ddthetalist + c(thetalist,dthetalist) + + g(thetalist) + Jtr(thetalist)Ftip + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + dthetalist = _as_batch_vec(dthetalist, thetalist.shape[-1]) + ddthetalist = _as_batch_vec(ddthetalist, thetalist.shape[-1]) + if g.ndim == 1: + g = g.unsqueeze(0).expand(thetalist.shape[0], 3) + Ftip = _as_batch_vec(Ftip, 6) + Mlist = Mlist if Mlist.ndim == 4 else Mlist.unsqueeze(0) + Glist = Glist if Glist.ndim == 4 else Glist.unsqueeze(0) + Slist = Slist if Slist.ndim == 3 else Slist.unsqueeze(0) + n = thetalist.shape[-1] + B = thetalist.shape[0] + taulist = torch.zeros((B, n), device=thetalist.device, dtype=thetalist.dtype) + for b in range(B): + Mi = torch.eye(4, device=thetalist.device, dtype=thetalist.dtype) + Ai = torch.zeros((6, n), device=thetalist.device, dtype=thetalist.dtype) + AdTi = [None] * (n + 1) + Vi = torch.zeros((6, n + 1), device=thetalist.device, dtype=thetalist.dtype) + Vdi = torch.zeros((6, n + 1), device=thetalist.device, dtype=thetalist.dtype) + Vdi[:, 0] = torch.cat([torch.zeros(3, device=thetalist.device, dtype=thetalist.dtype), -g[b]]) + AdTi[n] = Adjoint(TransInv(Mlist[b, n].unsqueeze(0))).squeeze(0) + Fi = Ftip[b].clone() + for i in range(n): + Mi = torch.matmul(Mi, Mlist[b, i]) + Ai[:, i] = torch.matmul(Adjoint(TransInv(Mi.unsqueeze(0))).squeeze(0), Slist[b, :, i]) + AdTi[i] = Adjoint(torch.matmul(MatrixExp6(VecTose3((Ai[:, i] * (-thetalist[b, i])).unsqueeze(0))).squeeze(0), + TransInv(Mlist[b, i].unsqueeze(0)).squeeze(0)).unsqueeze(0)).squeeze(0) + Vi[:, i + 1] = torch.matmul(AdTi[i], Vi[:, i]) + Ai[:, i] * dthetalist[b, i] + Vdi[:, i + 1] = torch.matmul(AdTi[i], Vdi[:, i]) + Ai[:, i] * ddthetalist[b, i] + torch.matmul(ad(Vi[:, i + 1].unsqueeze(0)).squeeze(0), Ai[:, i]) * dthetalist[b, i] + for i in range(n - 1, -1, -1): + Fi = torch.matmul(torch.transpose(AdTi[i + 1], -2, -1), Fi) + torch.matmul(Glist[b, i], Vdi[:, i + 1]) - torch.matmul(torch.transpose(ad(Vi[:, i + 1].unsqueeze(0)).squeeze(0), -2, -1), torch.matmul(Glist[b, i], Vi[:, i + 1])) + taulist[b, i] = torch.dot(Fi, Ai[:, i]) + return taulist + +def MassMatrix(thetalist: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes the mass matrix of an open chain robot based on configuration + + :param thetalist: A list of joint variables + :param Mlist: List of link frames i relative to i-1 at the home position + :param Glist: Spatial inertia matrices Gi of the links + :param Slist: Screw axes Si of the joints in a space frame, with axes as columns + :return: The numerical inertia matrix M(thetalist) + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + n = thetalist.shape[-1] + M = torch.zeros((thetalist.shape[0], n, n), device=thetalist.device, dtype=thetalist.dtype) + for i in range(n): + dd = torch.zeros_like(thetalist) + dd[:, i] = 1 + M[:, :, i] = InverseDynamics(thetalist, torch.zeros_like(thetalist), dd, torch.zeros((thetalist.shape[0], 3), device=thetalist.device, dtype=thetalist.dtype), torch.zeros((thetalist.shape[0], 6), device=thetalist.device, dtype=thetalist.dtype), Mlist, Glist, Slist) + return M + +def VelQuadraticForces(thetalist: torch.Tensor, dthetalist: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes Coriolis and centripetal terms in inverse dynamics. + + :param thetalist: Joint variables + :param dthetalist: Joint rates + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :return: c(thetalist, dthetalist) + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + dthetalist = _as_batch_vec(dthetalist, thetalist.shape[-1]) + return InverseDynamics(thetalist, dthetalist, torch.zeros_like(thetalist), torch.zeros((thetalist.shape[0], 3), device=thetalist.device, dtype=thetalist.dtype), torch.zeros((thetalist.shape[0], 6), device=thetalist.device, dtype=thetalist.dtype), Mlist, Glist, Slist) + +def GravityForces(thetalist: torch.Tensor, g: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes joint forces/torques required to overcome gravity. + + :param thetalist: Joint variables + :param g: Gravity vector + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :return: Gravity torque vector + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + if g.ndim == 1: + g = g.unsqueeze(0).expand(thetalist.shape[0], 3) + return InverseDynamics(thetalist, torch.zeros_like(thetalist), torch.zeros_like(thetalist), g, torch.zeros((thetalist.shape[0], 6), device=thetalist.device, dtype=thetalist.dtype), Mlist, Glist, Slist) + +def EndEffectorForces(thetalist: torch.Tensor, Ftip: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes joint forces/torques required only to create end-effector force. + + :param thetalist: Joint variables + :param Ftip: End-effector wrench in frame {n+1} + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :return: Joint forces/torques due to Ftip + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + Ftip = _as_batch_vec(Ftip, 6) + return InverseDynamics(thetalist, torch.zeros_like(thetalist), torch.zeros_like(thetalist), torch.zeros((thetalist.shape[0], 3), device=thetalist.device, dtype=thetalist.dtype), Ftip, Mlist, Glist, Slist) + +def ForwardDynamics(thetalist: torch.Tensor, dthetalist: torch.Tensor, taulist: torch.Tensor, g: torch.Tensor, Ftip: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Computes forward dynamics in the space frame for an open chain robot. + + :param thetalist: Joint variables + :param dthetalist: Joint rates + :param taulist: Joint forces/torques + :param g: Gravity vector + :param Ftip: End-effector wrench + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :return: Joint accelerations + """ + thetalist = _as_batch_vec(thetalist, Slist.shape[-1] if Slist.ndim > 1 else thetalist.shape[-1]) + dthetalist = _as_batch_vec(dthetalist, thetalist.shape[-1]) + taulist = _as_batch_vec(taulist, thetalist.shape[-1]) + if g.ndim == 1: + g = g.unsqueeze(0).expand(thetalist.shape[0], 3) + Ftip = _as_batch_vec(Ftip, 6) + M = MassMatrix(thetalist, Mlist, Glist, Slist) + rhs = taulist - VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist) - GravityForces(thetalist, g, Mlist, Glist, Slist) - EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist) + return torch.linalg.solve(M, rhs.unsqueeze(-1)).squeeze(-1) + +def EulerStep(thetalist: torch.Tensor, dthetalist: torch.Tensor, ddthetalist: torch.Tensor, dt: float): + """Performs a single Euler step for the given joint angles, velocities, and accelerations. + + :param thetalist: Joint angles, shape (N, 6) + :param dthetalist: Joint velocities, shape (N, 6) + :param ddthetalist: Joint accelerations, shape (N, 6) + :param dt: Time step for the Euler integration + :return: Updated joint angles and velocities after the Euler step + """ + return thetalist + dthetalist * dt, dthetalist + ddthetalist * dt + +def InverseDynamicsTrajectory(thetamat: torch.Tensor, dthetamatrix: torch.Tensor, ddthetamatrix: torch.Tensor, g: torch.Tensor, Ftipmatrix: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + """Calculates required joint torques along a trajectory using inverse dynamics. + + :param thetamat: N x n matrix of joint variables + :param dthetamatrix: N x n matrix of joint velocities + :param ddthetamatrix: N x n matrix of joint accelerations + :param g: Gravity vector + :param Ftipmatrix: N x 6 matrix of end-effector wrenches + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :return: N x n matrix of joint torques + """ + thetamat = _as_batch_mat(thetamat, thetamat.shape[-2], thetamat.shape[-1]) if thetamat.ndim == 2 else thetamat + if thetamat.ndim == 2: + thetamat = thetamat.unsqueeze(0) + if dthetamatrix.ndim == 2: + dthetamatrix = dthetamatrix.unsqueeze(0) + if ddthetamatrix.ndim == 2: + ddthetamatrix = ddthetamatrix.unsqueeze(0) + if Ftipmatrix.ndim == 2: + Ftipmatrix = Ftipmatrix.unsqueeze(0) + B, N, _ = thetamat.shape + out = torch.zeros_like(thetamat) + for b in range(B): + for i in range(N): + out[b, i] = InverseDynamics(thetamat[b, i], dthetamatrix[b, i], ddthetamatrix[b, i], g[b] if g.ndim == 2 else g, Ftipmatrix[b, i], Mlist[b] if Mlist.ndim == 4 else Mlist, Glist[b] if Glist.ndim == 4 else Glist, Slist[b] if Slist.ndim == 3 else Slist).squeeze(0) + return out + +def ForwardDynamicsTrajectory(thetamat: torch.Tensor, dthetamatrix: torch.Tensor, taumat: torch.Tensor, g: torch.Tensor, Ftipmatrix: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor, dt: float, intRes: int): + """Simulates motion given an open-loop history of joint torques. + + :param thetamat: Initial joint variables + :param dthetamatrix: Initial joint rates + :param taumat: N x n matrix of torques + :param g: Gravity vector + :param Ftipmatrix: N x 6 matrix of end-effector wrenches + :param Mlist: List of link frames + :param Glist: Spatial inertia matrices + :param Slist: Joint screw axes + :param dt: Time step between consecutive torque commands + :param intRes: Euler integration resolution per step + :return thetamat: N x n matrix of resulting joint angles + :return dthetamat: N x n matrix of resulting joint velocities + """ + if thetamat.ndim == 1: + thetamat = thetamat.unsqueeze(0) + if dthetamatrix.ndim == 1: + dthetamatrix = dthetamatrix.unsqueeze(0) + if taumat.ndim == 2: + taumat = taumat.unsqueeze(0) + if Ftipmatrix.ndim == 2: + Ftipmatrix = Ftipmatrix.unsqueeze(0) + B, N, n = taumat.shape + theta_out = torch.zeros((B, N, n), device=taumat.device, dtype=taumat.dtype) + dtheta_out = torch.zeros((B, N, n), device=taumat.device, dtype=taumat.dtype) + thetalist = thetamat.clone() + dthetalist = dthetamatrix.clone() + theta_out[:, 0] = thetalist + dtheta_out[:, 0] = dthetalist + for i in range(N - 1): + for _ in range(intRes): + dd = ForwardDynamics(thetalist, dthetalist, taumat[:, i], g, Ftipmatrix[:, i], Mlist, Glist, Slist) + thetalist, dthetalist = EulerStep(thetalist, dthetalist, dd, 1.0 * dt / intRes) + theta_out[:, i + 1] = thetalist + dtheta_out[:, i + 1] = dthetalist + return theta_out, dtheta_out + +def CubicTimeScaling(Tf, t): + """Computes s(t) for a cubic time scaling + + :param Tf: Total time of the motion in seconds from rest to rest + :param t: The current time t satisfying 0 < t < Tf + :return: The path parameter s(t) corresponding to a third-order + polynomial motion that begins and ends at zero velocity + + Example Input: + Tf = 2 + t = 0.6 + Output: + 0.216 + """ + return 3 * (1.0 * t / Tf) ** 2 - 2 * (1.0 * t / Tf) ** 3 + +def QuinticTimeScaling(Tf, t): + """Computes s(t) for a quintic time scaling + + :param Tf: Total time of the motion in seconds from rest to rest + :param t: The current time t satisfying 0 < t < Tf + :return: The path parameter s(t) corresponding to a fifth-order + polynomial motion that begins and ends at zero velocity and zero + acceleration + + Example Input: + Tf = 2 + t = 0.6 + Output: + 0.16308 + """ + return 10 * (1.0 * t / Tf) ** 3 - 15 * (1.0 * t / Tf) ** 4 \ + + 6 * (1.0 * t / Tf) ** 5 + +def JointTrajectory(thetalist: torch.Tensor, thetaend: torch.Tensor, Tf: float, N: int, method): + """Computes a straight-line trajectory in joint space. + + :param thetastart: Initial joint variables + :param thetaend: Final joint variables + :param Tf: Total time from rest to rest + :param N: Number of points in the trajectory + :param method: Time scaling method (3 for cubic, 5 for quintic) + :return: An N x n trajectory of joint variables + """ + thetalist = _as_batch_vec(thetalist, thetalist.shape[-1] if thetalist.ndim > 1 else thetalist.shape[0]) + thetaend = _as_batch_vec(thetaend, thetalist.shape[-1]) + N = int(N) + timegap = Tf / (N - 1.0) + traj = torch.zeros((thetalist.shape[0], N, thetalist.shape[-1]), device=thetalist.device, dtype=thetalist.dtype) + for i in range(N): + s = CubicTimeScaling(Tf, timegap * i) if method == 3 else QuinticTimeScaling(Tf, timegap * i) + traj[:, i] = s * thetaend + (1 - s) * thetalist + return traj + +def ScrewTrajectory(Xstart: torch.Tensor, Xend: torch.Tensor, Tf: float, N: int, method): + """Computes a trajectory as a list of N SE(3) matrices following screw motion. + + :param Xstart: Initial end-effector configuration + :param Xend: Final end-effector configuration + :param Tf: Total motion time + :param N: Number of points + :param method: Time scaling method (3 or 5) + :return: N SE(3) matrices, separated by Tf/(N-1) + """ + Xstart = _as_batch_mat(Xstart, 4, 4) + Xend = _as_batch_mat(Xend, 4, 4) + N = int(N) + timegap = Tf / (N - 1.0) + traj = torch.zeros((Xstart.shape[0], N, 4, 4), device=Xstart.device, dtype=Xstart.dtype) + for i in range(N): + s = CubicTimeScaling(Tf, timegap * i) if method == 3 else QuinticTimeScaling(Tf, timegap * i) + traj[:, i] = torch.bmm(Xstart, MatrixExp6(MatrixLog6(torch.bmm(TransInv(Xstart), Xend)) * s)) + return traj + +def CartesianTrajectory(Xstart: torch.Tensor, Xend: torch.Tensor, Tf: float, N: int, method): + """Computes a trajectory with straight-line translation and decoupled rotation. + + :param Xstart: Initial end-effector configuration + :param Xend: Final end-effector configuration + :param Tf: Total motion time + :param N: Number of points + :param method: Time scaling method (3 or 5) + :return: N SE(3) matrices, separated by Tf/(N-1) + """ + Xstart = _as_batch_mat(Xstart, 4, 4) + Xend = _as_batch_mat(Xend, 4, 4) + N = int(N) + timegap = Tf / (N - 1.0) + traj = torch.zeros((Xstart.shape[0], N, 4, 4), device=Xstart.device, dtype=Xstart.dtype) + Rstart, pstart = TransToRp(Xstart) + Rend, pend = TransToRp(Xend) + for i in range(N): + s = CubicTimeScaling(Tf, timegap * i) if method == 3 else QuinticTimeScaling(Tf, timegap * i) + R = torch.bmm(Rstart, MatrixExp3(MatrixLog3(torch.bmm(torch.transpose(Rstart, -2, -1), Rend)) * s)) + p = s * pend + (1 - s) * pstart + traj[:, i] = RpToTrans(R, p.squeeze(-1)) + return traj + +def ComputedTorque(**kwargs): + """Computes joint control torques at a particular time instant. + + Keyword args follow the NumPy implementation: + thetalist, dthetalist, eint, g, Mlist, Glist, Slist, + thetalistd, dthetalistd, ddthetalistd, Kp, Ki, Kd. + """ + thetalist = kwargs["thetalist"] + dthetalist = kwargs["dthetalist"] + eint = kwargs["eint"] + g = kwargs["g"] + Mlist = kwargs["Mlist"] + Glist = kwargs["Glist"] + Slist = kwargs["Slist"] + thetalistd = kwargs["thetalistd"] + dthetalistd = kwargs["dthetalistd"] + ddthetalistd = kwargs["ddthetalistd"] + Kp = kwargs["Kp"] + Ki = kwargs["Ki"] + Kd = kwargs["Kd"] + thetalist = _as_batch_vec(thetalist, thetalist.shape[-1] if thetalist.ndim > 1 else thetalist.shape[0]) + dthetalist = _as_batch_vec(dthetalist, thetalist.shape[-1]) + eint = _as_batch_vec(eint, thetalist.shape[-1]) + thetalistd = _as_batch_vec(thetalistd, thetalist.shape[-1]) + dthetalistd = _as_batch_vec(dthetalistd, thetalist.shape[-1]) + ddthetalistd = _as_batch_vec(ddthetalistd, thetalist.shape[-1]) + e = thetalistd - thetalist + return torch.bmm(MassMatrix(thetalist, Mlist, Glist, Slist), (Kp * e + Ki * (eint + e) + Kd * (dthetalistd - dthetalist)).unsqueeze(-1)).squeeze(-1) + InverseDynamics(thetalist, dthetalist, ddthetalistd, g, torch.zeros((thetalist.shape[0], 6), device=thetalist.device, dtype=thetalist.dtype), Mlist, Glist, Slist) + +def SimulateControl(**kwargs): + """Simulates computed torque control over a desired trajectory. + + Keyword args follow the NumPy implementation: + thetalist, dthetalist, g, Ftipmat, Mlist, Glist, Slist, + thetamatd, dthetamatd, ddthetamatd, gtilde, Mtildelist, Gtildelist, + Kp, Ki, Kd, dt, intRes. + :return taumat: Commanded joint torques over time + :return thetamat: Actual joint angles over time + """ + thetalist = kwargs["thetalist"] + dthetalist = kwargs["dthetalist"] + g = kwargs["g"] + Ftipmat = kwargs["Ftipmat"] + Mlist = kwargs["Mlist"] + Glist = kwargs["Glist"] + Slist = kwargs["Slist"] + thetamatd = kwargs["thetamatd"] + dthetamatd = kwargs["dthetamatd"] + ddthetamatd = kwargs["ddthetamatd"] + gtilde = kwargs["gtilde"] + Mtildelist = kwargs["Mtildelist"] + Gtildelist = kwargs["Gtildelist"] + Kp = kwargs["Kp"] + Ki = kwargs["Ki"] + Kd = kwargs["Kd"] + dt = kwargs["dt"] + intRes = kwargs["intRes"] + if thetalist.ndim == 1: + thetalist = thetalist.unsqueeze(0) + if dthetalist.ndim == 1: + dthetalist = dthetalist.unsqueeze(0) + if thetamatd.ndim == 2: + thetamatd = thetamatd.unsqueeze(0) + if dthetamatd.ndim == 2: + dthetamatd = dthetamatd.unsqueeze(0) + if ddthetamatd.ndim == 2: + ddthetamatd = ddthetamatd.unsqueeze(0) + if Ftipmat.ndim == 2: + Ftipmat = Ftipmat.unsqueeze(0) + B, N, n = thetamatd.shape + thetacurrent = thetalist.clone() + dthetacurrent = dthetalist.clone() + eint = torch.zeros((B, n), device=thetamatd.device, dtype=thetamatd.dtype) + taumat = torch.zeros((B, N, n), device=thetamatd.device, dtype=thetamatd.dtype) + thetamat = torch.zeros((B, N, n), device=thetamatd.device, dtype=thetamatd.dtype) + for i in range(N): + taulist = ComputedTorque( + thetalist=thetacurrent, dthetalist=dthetacurrent, eint=eint, g=gtilde, + Mlist=Mtildelist, Glist=Gtildelist, Slist=Slist, thetalistd=thetamatd[:, i], + dthetalistd=dthetamatd[:, i], ddthetalistd=ddthetamatd[:, i], Kp=Kp, Ki=Ki, Kd=Kd + ) + for _ in range(intRes): + dd = ForwardDynamics(thetacurrent, dthetacurrent, taulist, g, Ftipmat[:, i], Mlist, Glist, Slist) + thetacurrent, dthetacurrent = EulerStep(thetacurrent, dthetacurrent, dd, 1.0 * dt / intRes) + taumat[:, i] = taulist + thetamat[:, i] = thetacurrent + eint = eint + dt * (thetamatd[:, i] - thetacurrent) + return taumat, thetamat diff --git a/packages/PyTorch/tests/README.md b/packages/PyTorch/tests/README.md new file mode 100644 index 0000000..963d857 --- /dev/null +++ b/packages/PyTorch/tests/README.md @@ -0,0 +1,26 @@ +# Unit tests for `pytorch_mr` + +Tests live under `packages/PyTorch/tests/` (sibling of `src/`), matching a standard PyPA layout. + +They compare outputs to the NumPy reference `modern_robotics` in `packages/Python`. `conftest.py` adds `packages/Python` to `sys.path` when `modern_robotics` is not installed. + +**Recommended (from `packages/PyTorch`):** + +```bash +python -m pytest tests -q +``` + +**From `packages/`:** + +```bash +PYTHONPATH=PyTorch/src:Python python -m pytest PyTorch/tests -q +``` + +**Chapter summary:** + +```bash +cd packages +PYTHONPATH=PyTorch/src:Python python PyTorch/tests/run_tests_by_chapter.py +``` + +With `pip install modern_robotics` (or editable `../Python`), you only need `src` on the path for an uninstalled checkout; after `pip install -e .`, `pytorch_mr` resolves without extra `PYTHONPATH`. diff --git a/packages/PyTorch/tests/conftest.py b/packages/PyTorch/tests/conftest.py new file mode 100644 index 0000000..a2848ec --- /dev/null +++ b/packages/PyTorch/tests/conftest.py @@ -0,0 +1,134 @@ +import os +import sys +from typing import Any, Callable, Sequence + +import numpy as np +import pytest +import torch + +_HERE = os.path.dirname(__file__) + +# Installed package: pytorch_mr is on sys.path. Local checkout: add src/. +try: + import modern_robotics.core as mr_np # noqa: E402 +except ImportError: + _PKG_PYTHON_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..", "Python")) + if _PKG_PYTHON_ROOT not in sys.path: + sys.path.insert(0, _PKG_PYTHON_ROOT) + import modern_robotics.core as mr_np # noqa: E402 + +try: + import pytorch_mr.core as mr_t # noqa: E402 +except ImportError: + _SRC_ROOT = os.path.abspath(os.path.join(_HERE, "..", "src")) + if _SRC_ROOT not in sys.path: + sys.path.insert(0, _SRC_ROOT) + import pytorch_mr.core as mr_t # noqa: E402 + + +@pytest.fixture(scope="session") +def np_ref(): + return mr_np + + +@pytest.fixture(scope="session") +def pt_impl(): + return mr_t + + +def to_numpy(x: Any) -> Any: + if isinstance(x, torch.Tensor): + return x.detach().cpu().numpy() + return x + + +def to_tensor(x: Any, *, dtype=None, device=None) -> torch.Tensor: + if isinstance(x, torch.Tensor): + t = x + else: + t = torch.tensor(x) + if dtype is not None: + t = t.to(dtype=dtype) + if device is not None: + t = t.to(device=device) + return t + + +def _iter_batch(x: torch.Tensor) -> Sequence[torch.Tensor]: + if x.ndim == 0: + return [x] + if x.ndim == 1: + return [x] + if x.ndim >= 2: + return [x[i] for i in range(x.shape[0])] + return [x] + + +def assert_torch_matches_numpy( + torch_out: Any, + numpy_out: Any, + *, + atol=1e-6, + rtol=1e-6, +): + if isinstance(torch_out, tuple): + assert isinstance(numpy_out, tuple) + assert len(torch_out) == len(numpy_out) + for a, b in zip(torch_out, numpy_out): + assert_torch_matches_numpy(a, b, atol=atol, rtol=rtol) + return + + if isinstance(torch_out, torch.Tensor): + np_t = torch_out.detach().cpu().numpy() + np_n = np.array(numpy_out) + assert np_t.shape == np_n.shape + assert np.allclose(np_t, np_n, atol=atol, rtol=rtol), (np_t, np_n) + return + + assert torch_out == numpy_out + + +def apply_numpy_per_batch( + f_np: Callable[..., Any], + *args_t: Any, + squeeze_last_dim1: bool = True, +) -> Any: + """ + Call NumPy reference function `f_np` on each batch element. + + For tensor args: + - If shape is (N, ..., 1) and squeeze_last_dim1=True, squeeze the last dim + before passing to NumPy (to align with reference signatures). + - If shape is (N, ...), iterate over axis 0. + - If shape is (...), treat as single item. + """ + batch = None + for a in args_t: + if isinstance(a, torch.Tensor) and a.ndim >= 2: + batch = a.shape[0] + break + if batch is None: + np_args = [] + for a in args_t: + if isinstance(a, torch.Tensor): + x = a + if squeeze_last_dim1 and x.ndim >= 1 and x.shape[-1] == 1: + x = x.squeeze(-1) + np_args.append(to_numpy(x)) + else: + np_args.append(a) + return f_np(*np_args) + + outs = [] + for i in range(batch): + np_args_i = [] + for a in args_t: + if isinstance(a, torch.Tensor): + x = a[i] if a.ndim >= 2 else a + if squeeze_last_dim1 and x.ndim >= 1 and x.shape[-1] == 1: + x = x.squeeze(-1) + np_args_i.append(to_numpy(x)) + else: + np_args_i.append(a) + outs.append(f_np(*np_args_i)) + return outs diff --git a/packages/PyTorch/tests/run_tests_by_chapter.py b/packages/PyTorch/tests/run_tests_by_chapter.py new file mode 100644 index 0000000..b4d2635 --- /dev/null +++ b/packages/PyTorch/tests/run_tests_by_chapter.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python3 +import os +import subprocess +import sys +from dataclasses import dataclass + + +@dataclass +class ChapterRun: + name: str + path: str + collected: int + exit_code: int + + +def _collect_count(pytest_args, env, cwd): + p = subprocess.run( + [sys.executable, "-m", "pytest", "--collect-only", "-q", *pytest_args], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, + cwd=cwd, + ) + lines = [ln.strip() for ln in p.stdout.splitlines() if ln.strip()] + items = [ln for ln in lines if "::" in ln and not ln.startswith("<") and "warnings summary" not in ln.lower()] + return len(items) + + +def _run(pytest_args, env, cwd): + p = subprocess.run( + [sys.executable, "-m", "pytest", "-q", *pytest_args], + env=env, + cwd=cwd, + ) + return p.returncode + + +def main(): + here = os.path.abspath(os.path.dirname(__file__)) + # tests/ -> PyTorch/ -> packages/ + packages_dir = os.path.abspath(os.path.join(here, "..", "..")) + torch_src = os.path.join(packages_dir, "PyTorch", "src") + python_pkg_root = os.path.join(packages_dir, "Python") + env = dict(os.environ) + sep = os.pathsep + extra = torch_src + sep + python_pkg_root + env["PYTHONPATH"] = extra + (sep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "") + + chapters = [ + ("ch1_3", os.path.join(here, "test_ch1_3_numpy_parity.py")), + ("ch3", os.path.join(here, "test_ch3.py")), + ("ch4_6", os.path.join(here, "test_ch4_6.py")), + ("ch8", os.path.join(here, "test_ch8.py")), + ("ch9_11", os.path.join(here, "test_ch9_11.py")), + ] + + runs = [] + total_collected = 0 + any_failed = False + + for name, path in chapters: + rel = os.path.relpath(path, packages_dir) + collected = _collect_count([rel], env, cwd=packages_dir) + exit_code = _run([rel], env, cwd=packages_dir) + runs.append(ChapterRun(name=name, path=rel, collected=collected, exit_code=exit_code)) + total_collected += collected + any_failed = any_failed or (exit_code != 0) + + print("\n=== Chapter test summary ===") + for r in runs: + status = "PASS" if r.exit_code == 0 else "FAIL" + print(f"{r.name:>6} | {status} | collected={r.collected:4d} | {r.path}") + print(f"TOTAL | {'PASS' if not any_failed else 'FAIL'} | collected={total_collected:4d}") + + return 1 if any_failed else 0 + + +if __name__ == "__main__": + raise SystemExit(main()) diff --git a/packages/PyTorch/tests/test_ch1_3_numpy_parity.py b/packages/PyTorch/tests/test_ch1_3_numpy_parity.py new file mode 100644 index 0000000..c1f7334 --- /dev/null +++ b/packages/PyTorch/tests/test_ch1_3_numpy_parity.py @@ -0,0 +1,170 @@ +import numpy as np +import pytest +import torch + + +@pytest.mark.parametrize( + "shape", + [ + (3,), # (dim,) + (1, 3), # (N, dim) + (5, 3), # (N, dim) + (5, 3, 1), # (N, dim, 1) + ], +) +def test_normalize_numpy_parity(pt_impl, np_ref, shape): + x = torch.randn(*shape, dtype=torch.float64) + y = pt_impl.Normalize(x) + + # NumPy Normalize expects a vector, not batched; apply per-batch + if x.ndim == 1: + y_np = np_ref.Normalize(x.detach().cpu().numpy()) + assert y.shape == (1, 3, 1) # current PyTorch Normalize contract + assert np.allclose(y.squeeze(0).squeeze(-1).cpu().numpy(), y_np, atol=1e-6) + else: + batch = x.shape[0] + y_np = [] + for i in range(batch): + xi = x[i] + if xi.ndim == 2 and xi.shape[-1] == 1: + xi = xi.squeeze(-1) + y_np.append(np_ref.Normalize(xi.detach().cpu().numpy())) + y_np = np.stack(y_np, axis=0) + assert np.allclose(y.squeeze(-1).cpu().numpy(), y_np, atol=1e-6) + + +@pytest.mark.parametrize("val", [-1e-7, 0.0, 1e-5, 1e-4]) +def test_nearzero_numpy_parity(pt_impl, np_ref, val): + t = torch.tensor(val, dtype=torch.float64) + out = bool(pt_impl.NearZero(t)) + expected = bool(np_ref.NearZero(float(val))) + assert out == expected + + +def test_rotinv_numpy_parity(pt_impl, np_ref): + R = torch.tensor([[0.0, 0.0, 1.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0]], dtype=torch.float64) + out = pt_impl.RotInv(R).squeeze(0).cpu().numpy() + expected = np_ref.RotInv(R.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-12) + + +@pytest.mark.parametrize( + "omega_shape", + [(3,), (1, 3), (7, 3), (7, 3, 1)], +) +def test_vectoso3_and_back_numpy_parity(pt_impl, np_ref, omega_shape): + omega = torch.randn(*omega_shape, dtype=torch.float64) + so3 = pt_impl.VecToso3(omega) + + if omega.ndim == 1: + expected = np_ref.VecToso3(omega.cpu().numpy()) + assert np.allclose(so3.squeeze(0).cpu().numpy(), expected, atol=1e-12) + back = pt_impl.so3ToVec(so3).squeeze(0).squeeze(-1).cpu().numpy() + back_expected = np_ref.so3ToVec(expected) + assert np.allclose(back, back_expected, atol=1e-12) + else: + batch = omega.shape[0] + expected = [] + for i in range(batch): + oi = omega[i] + if oi.ndim == 2 and oi.shape[-1] == 1: + oi = oi.squeeze(-1) + expected.append(np_ref.VecToso3(oi.cpu().numpy())) + expected = np.stack(expected, axis=0) + assert np.allclose(so3.cpu().numpy(), expected, atol=1e-12) + + +def test_axisang3_numpy_parity(pt_impl, np_ref): + expc3 = torch.tensor([1.0, 2.0, 3.0], dtype=torch.float64) + omghat_t, theta_t = pt_impl.AxisAng3(expc3) + + omghat_np, theta_np = np_ref.AxisAng3(expc3.cpu().numpy()) + assert np.allclose(omghat_t.squeeze(0).squeeze(-1).cpu().numpy(), omghat_np, atol=1e-6) + assert np.allclose(theta_t.squeeze(0).cpu().numpy(), theta_np, atol=1e-6) + + +def test_matrixexp3_numpy_example_parity(pt_impl, np_ref): + so3 = torch.tensor([[0.0, -3.0, 2.0], + [3.0, 0.0, -1.0], + [-2.0, 1.0, 0.0]], dtype=torch.float64) + out = pt_impl.MatrixExp3(so3).squeeze(0).cpu().numpy() + expected = np_ref.MatrixExp3(so3.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-6) + + +def test_matrixlog3_numpy_example_parity(pt_impl, np_ref): + R = torch.tensor([[0.0, 0.0, 1.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0]], dtype=torch.float64) + out = pt_impl.MatrixLog3(R).squeeze(0).cpu().numpy() + expected = np_ref.MatrixLog3(R.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-6) + + +def test_rp_trans_roundtrip_numpy_parity(pt_impl, np_ref): + R = torch.tensor([[1.0, 0.0, 0.0], + [0.0, 0.0, -1.0], + [0.0, 1.0, 0.0]], dtype=torch.float64) + p = torch.tensor([1.0, 2.0, 5.0], dtype=torch.float64) + Tt = pt_impl.RpToTrans(R.unsqueeze(0), p.unsqueeze(0)).squeeze(0).cpu().numpy() + Tn = np_ref.RpToTrans(R.cpu().numpy(), p.cpu().numpy()) + assert np.allclose(Tt, Tn, atol=1e-12) + + Rt, pt = pt_impl.TransToRp(torch.tensor(Tn, dtype=torch.float64)) # type: ignore + Rt = Rt.squeeze(0).cpu().numpy() + pt = pt.squeeze(0).squeeze(-1).cpu().numpy() + Rn, pn = np_ref.TransToRp(Tn) + assert np.allclose(Rt, Rn, atol=1e-12) + assert np.allclose(pt, pn, atol=1e-12) + + +def test_transinv_numpy_example_parity(pt_impl, np_ref): + T = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 1.0, 0.0, 3.0], + [0.0, 0.0, 0.0, 1.0]], dtype=torch.float64) + out = pt_impl.TransInv(T).squeeze(0).cpu().numpy() + expected = np_ref.TransInv(T.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-12) + + +def test_vectose3_and_back_numpy_parity(pt_impl, np_ref): + V = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0], dtype=torch.float64) + se3 = pt_impl.VecTose3(V.unsqueeze(0)).squeeze(0).cpu().numpy() + expected = np_ref.VecTose3(V.cpu().numpy()) + assert np.allclose(se3, expected, atol=1e-12) + + back = pt_impl.se3ToVec(torch.tensor(expected, dtype=torch.float64)).squeeze(0).squeeze(-1).cpu().numpy() + back_expected = np_ref.se3ToVec(expected) + assert np.allclose(back, back_expected, atol=1e-12) + + +def test_adjoint_numpy_example_parity(pt_impl, np_ref): + T = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 1.0, 0.0, 3.0], + [0.0, 0.0, 0.0, 1.0]], dtype=torch.float64) + out = pt_impl.Adjoint(T).squeeze(0).cpu().numpy() + expected = np_ref.Adjoint(T.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-12) + + +def test_matrixexp6_and_log6_numpy_examples_parity(pt_impl, np_ref): + se3 = torch.tensor([[0.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.57079632, 2.35619449], + [0.0, 1.57079632, 0.0, 2.35619449], + [0.0, 0.0, 0.0, 0.0]], dtype=torch.float64) + out = pt_impl.MatrixExp6(se3.unsqueeze(0)).squeeze(0).cpu().numpy() + expected = np_ref.MatrixExp6(se3.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-6) + + T = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, -1.0, 0.0], + [0.0, 1.0, 0.0, 3.0], + [0.0, 0.0, 0.0, 1.0]], dtype=torch.float64) + out_log = pt_impl.MatrixLog6(T).squeeze(0).cpu().numpy() + expected_log = np_ref.MatrixLog6(T.cpu().numpy()) + assert np.allclose(out_log, expected_log, atol=1e-6) + diff --git a/packages/PyTorch/tests/test_ch3.py b/packages/PyTorch/tests/test_ch3.py new file mode 100644 index 0000000..165a34e --- /dev/null +++ b/packages/PyTorch/tests/test_ch3.py @@ -0,0 +1,67 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import pytorch_mr.core as pmr + + +def test_screw_to_axis_example(): + q = torch.tensor([3.0, 0.0, 0.0]) + s = torch.tensor([0.0, 0.0, 1.0]) + h = torch.tensor(2.0) + out = pmr.ScrewToAxis(q, s, h) + expected = np_mr.ScrewToAxis(q.numpy(), s.numpy(), float(h)) + assert np.allclose(out.cpu().numpy(), expected[None, :], atol=1e-6) + + +def test_axisang6_example(): + expc6 = torch.tensor([1.0, 0.0, 0.0, 1.0, 2.0, 3.0]) + S, theta = pmr.AxisAng6(expc6) + S_np, theta_np = np_mr.AxisAng6(expc6.numpy()) + assert np.allclose(S.squeeze(0).cpu().numpy(), S_np, atol=1e-6) + assert np.allclose(theta.squeeze(0).cpu().numpy(), theta_np, atol=1e-6) + + +def test_project_to_so3_example(): + mat = torch.tensor([[0.675, 0.150, 0.720], + [0.370, 0.771, -0.511], + [-0.630, 0.619, 0.472]], dtype=torch.float64) + out = pmr.ProjectToSO3(mat).squeeze(0) + expected = np_mr.ProjectToSO3(mat.cpu().numpy()) + assert np.allclose(out.cpu().numpy(), expected, atol=1e-5) + + +def test_project_to_se3_example(): + mat = torch.tensor([[0.675, 0.150, 0.720, 1.2], + [0.370, 0.771, -0.511, 5.4], + [-0.630, 0.619, 0.472, 3.6], + [0.003, 0.002, 0.010, 0.9]], dtype=torch.float64) + out = pmr.ProjectToSE3(mat).squeeze(0) + expected = np_mr.ProjectToSE3(mat.cpu().numpy()) + assert np.allclose(out.cpu().numpy(), expected, atol=1e-6) + + +def test_distance_and_membership_examples(): + so3_bad = torch.tensor([[1.0, 0.0, 0.0], + [0.0, 0.1, -0.95], + [0.0, 1.0, 0.1]]) + dso3 = pmr.DistanceToSO3(so3_bad).squeeze(0).cpu().numpy() + expected_dso3 = np_mr.DistanceToSO3(so3_bad.cpu().numpy()) + assert np.allclose(dso3, expected_dso3, atol=1e-5) + assert bool(pmr.TestIfSO3(so3_bad).squeeze(0)) == bool(np_mr.TestIfSO3(so3_bad.cpu().numpy())) + + se3_bad = torch.tensor([[1.0, 0.0, 0.0, 1.2], + [0.0, 0.1, -0.95, 1.5], + [0.0, 1.0, 0.1, -0.9], + [0.0, 0.0, 0.1, 0.98]]) + dse3 = pmr.DistanceToSE3(se3_bad).squeeze(0).cpu().numpy() + expected_dse3 = np_mr.DistanceToSE3(se3_bad.cpu().numpy()) + assert np.allclose(dse3, expected_dse3, atol=1e-5) + assert bool(pmr.TestIfSE3(se3_bad).squeeze(0)) == bool(np_mr.TestIfSE3(se3_bad.cpu().numpy())) + + +def test_broadcast_and_device_basic(): + x = torch.tensor([[1.0, 2.0, 3.0], [0.0, 1.0, 0.0]]) + out = pmr.Normalize(x) + assert out.shape[-2:] == (3, 1) + assert out.device == x.device diff --git a/packages/PyTorch/tests/test_ch4_6.py b/packages/PyTorch/tests/test_ch4_6.py new file mode 100644 index 0000000..1f6179a --- /dev/null +++ b/packages/PyTorch/tests/test_ch4_6.py @@ -0,0 +1,79 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import pytorch_mr.core as pmr + + +def test_fkin_body_example(): + M = torch.tensor([[-1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 6.0], + [0.0, 0.0, -1.0, 2.0], + [0.0, 0.0, 0.0, 1.0]]) + Blist = torch.tensor([[0.0, 0.0, -1.0, 2.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.1]]).T + thetalist = torch.tensor([torch.pi / 2.0, 3.0, torch.pi]) + out = pmr.FKinBody(M, Blist, thetalist).squeeze(0).cpu().numpy() + expected = np_mr.FKinBody(M.cpu().numpy(), Blist.cpu().numpy(), thetalist.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-5) + + +def test_fkin_space_example(): + M = torch.tensor([[-1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 6.0], + [0.0, 0.0, -1.0, 2.0], + [0.0, 0.0, 0.0, 1.0]]) + Slist = torch.tensor([[0.0, 0.0, 1.0, 4.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, -1.0, -6.0, 0.0, -0.1]]).T + thetalist = torch.tensor([torch.pi / 2.0, 3.0, torch.pi]) + out = pmr.FKinSpace(M, Slist, thetalist).squeeze(0).cpu().numpy() + expected = np_mr.FKinSpace(M.cpu().numpy(), Slist.cpu().numpy(), thetalist.cpu().numpy()) + assert np.allclose(out, expected, atol=1e-5) + + +def test_jacobians_shapes(): + Blist = torch.tensor([[0.0, 0.0, 1.0, 0.0, 0.2, 0.2], + [1.0, 0.0, 0.0, 2.0, 0.0, 3.0], + [0.0, 1.0, 0.0, 0.0, 2.0, 1.0], + [1.0, 0.0, 0.0, 0.2, 0.3, 0.4]]).T + thetalist = torch.tensor([0.2, 1.1, 0.1, 1.2]) + Jb = pmr.JacobianBody(Blist, thetalist) + Js = pmr.JacobianSpace(Blist, thetalist) + assert Jb.shape == (1, 6, 4) + assert Js.shape == (1, 6, 4) + + Jb_np = np_mr.JacobianBody(Blist.cpu().numpy(), thetalist.cpu().numpy()) + Js_np = np_mr.JacobianSpace(Blist.cpu().numpy(), thetalist.cpu().numpy()) + assert np.allclose(Jb.squeeze(0).cpu().numpy(), Jb_np, atol=1e-5) + assert np.allclose(Js.squeeze(0).cpu().numpy(), Js_np, atol=1e-5) + + +def test_ikin_body_and_space_converge(): + Blist = torch.tensor([[0.0, 0.0, -1.0, 2.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, 1.0, 0.0, 0.0, 0.1]]).T + Slist = torch.tensor([[0.0, 0.0, 1.0, 4.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 0.0, 1.0, 0.0], + [0.0, 0.0, -1.0, -6.0, 0.0, -0.1]]).T + M = torch.tensor([[-1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 6.0], + [0.0, 0.0, -1.0, 2.0], + [0.0, 0.0, 0.0, 1.0]]) + T = torch.tensor([[0.0, 1.0, 0.0, -5.0], + [1.0, 0.0, 0.0, 4.0], + [0.0, 0.0, -1.0, 1.6858], + [0.0, 0.0, 0.0, 1.0]]) + guess = torch.tensor([1.5, 2.5, 3.0]) + th_b, ok_b = pmr.IKinBody(Blist, M, T, guess, eomg=0.01, ev=0.001) + th_s, ok_s = pmr.IKinSpace(Slist, M, T, guess, eomg=0.01, ev=0.001) + assert bool(ok_b.squeeze(0)) + assert bool(ok_s.squeeze(0)) + + th_b_np, ok_b_np = np_mr.IKinBody(Blist.cpu().numpy(), M.cpu().numpy(), T.cpu().numpy(), guess.cpu().numpy(), 0.01, 0.001) + th_s_np, ok_s_np = np_mr.IKinSpace(Slist.cpu().numpy(), M.cpu().numpy(), T.cpu().numpy(), guess.cpu().numpy(), 0.01, 0.001) + assert bool(ok_b_np) + assert bool(ok_s_np) + assert np.allclose(th_b.squeeze(0).cpu().numpy(), th_b_np, atol=1e-3) + assert np.allclose(th_s.squeeze(0).cpu().numpy(), th_s_np, atol=1e-3) diff --git a/packages/PyTorch/tests/test_ch8.py b/packages/PyTorch/tests/test_ch8.py new file mode 100644 index 0000000..bfe7d4c --- /dev/null +++ b/packages/PyTorch/tests/test_ch8.py @@ -0,0 +1,105 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import pytorch_mr.core as pmr + + +def robot3(): + M01 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.089159], + [0.0, 0.0, 0.0, 1.0]]) + M12 = torch.tensor([[0.0, 0.0, 1.0, 0.28], + [0.0, 1.0, 0.0, 0.13585], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0]]) + M23 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.1197], + [0.0, 0.0, 1.0, 0.395], + [0.0, 0.0, 0.0, 1.0]]) + M34 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.14225], + [0.0, 0.0, 0.0, 1.0]]) + G1 = torch.diag(torch.tensor([0.010267, 0.010267, 0.00666, 3.7, 3.7, 3.7])) + G2 = torch.diag(torch.tensor([0.22689, 0.22689, 0.0151074, 8.393, 8.393, 8.393])) + G3 = torch.diag(torch.tensor([0.0494433, 0.0494433, 0.004095, 2.275, 2.275, 2.275])) + Mlist = torch.stack([M01, M12, M23, M34], dim=0) + Glist = torch.stack([G1, G2, G3], dim=0) + Slist = torch.tensor([[1.0, 0.0, 1.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, -0.089, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.089, 0.0, 0.425]]).T + return Mlist, Glist, Slist + + +def test_ad_example(): + V = torch.tensor([1.0, 2.0, 3.0, 4.0, 5.0, 6.0]) + out = pmr.ad(V).squeeze(0) + expected = torch.tensor([[0.0, -3.0, 2.0, 0.0, 0.0, 0.0], + [3.0, 0.0, -1.0, 0.0, 0.0, 0.0], + [-2.0, 1.0, 0.0, 0.0, 0.0, 0.0], + [0.0, -6.0, 5.0, 0.0, -3.0, 2.0], + [6.0, 0.0, -4.0, 3.0, 0.0, -1.0], + [-5.0, 4.0, 0.0, -2.0, 1.0, 0.0]]) + assert torch.allclose(out, expected, atol=1e-6) + expected_np = np_mr.ad(V.cpu().numpy()) + assert np.allclose(out.cpu().numpy(), expected_np, atol=1e-6) + + +def test_inverse_dynamics_and_related_examples(): + Mlist, Glist, Slist = robot3() + thetalist = torch.tensor([0.1, 0.1, 0.1]) + dthetalist = torch.tensor([0.1, 0.2, 0.3]) + ddthetalist = torch.tensor([2.0, 1.5, 1.0]) + g = torch.tensor([0.0, 0.0, -9.8]) + Ftip = torch.tensor([1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + + tau = pmr.InverseDynamics(thetalist, dthetalist, ddthetalist, g, Ftip, Mlist, Glist, Slist).squeeze(0) + expected_tau = np_mr.InverseDynamics( + thetalist.cpu().numpy(), dthetalist.cpu().numpy(), ddthetalist.cpu().numpy(), + g.cpu().numpy(), Ftip.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy() + ) + assert np.allclose(tau.cpu().numpy(), expected_tau, atol=1e-3) + + M = pmr.MassMatrix(thetalist, Mlist, Glist, Slist).squeeze(0) + expected_M = np_mr.MassMatrix(thetalist.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy()) + assert np.allclose(M.cpu().numpy(), expected_M, atol=1e-3) + + c = pmr.VelQuadraticForces(thetalist, dthetalist, Mlist, Glist, Slist).squeeze(0) + expected_c = np_mr.VelQuadraticForces(thetalist.cpu().numpy(), dthetalist.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy()) + assert np.allclose(c.cpu().numpy(), expected_c, atol=1e-4) + + grav = pmr.GravityForces(thetalist, g, Mlist, Glist, Slist).squeeze(0) + expected_g = np_mr.GravityForces(thetalist.cpu().numpy(), g.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy()) + assert np.allclose(grav.cpu().numpy(), expected_g, atol=1e-3) + + tip = pmr.EndEffectorForces(thetalist, Ftip, Mlist, Glist, Slist).squeeze(0) + expected_tip = np_mr.EndEffectorForces(thetalist.cpu().numpy(), Ftip.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy()) + assert np.allclose(tip.cpu().numpy(), expected_tip, atol=1e-3) + + +def test_forward_dynamics_and_trajectory_shapes(): + Mlist, Glist, Slist = robot3() + thetalist = torch.tensor([0.1, 0.1, 0.1]) + dthetalist = torch.tensor([0.1, 0.2, 0.3]) + taulist = torch.tensor([0.5, 0.6, 0.7]) + g = torch.tensor([0.0, 0.0, -9.8]) + Ftip = torch.tensor([1.0, 1.0, 1.0, 1.0, 1.0, 1.0]) + dd = pmr.ForwardDynamics(thetalist, dthetalist, taulist, g, Ftip, Mlist, Glist, Slist).squeeze(0) + expected_dd = np_mr.ForwardDynamics( + thetalist.cpu().numpy(), dthetalist.cpu().numpy(), taulist.cpu().numpy(), + g.cpu().numpy(), Ftip.cpu().numpy(), Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy() + ) + assert np.allclose(dd.cpu().numpy(), expected_dd, atol=1e-2) + + thetamat = torch.stack([thetalist, thetalist + 0.01], dim=0) + dthetamat = torch.zeros_like(thetamat) + ddthetamat = torch.zeros_like(thetamat) + Ftipmat = torch.ones((2, 6)) + taumat = pmr.InverseDynamicsTrajectory(thetamat, dthetamat, ddthetamat, g, Ftipmat, Mlist, Glist, Slist) + assert taumat.shape == (1, 2, 3) + + th, dth = pmr.ForwardDynamicsTrajectory(thetalist, dthetalist, torch.ones((2, 3)), g, torch.ones((2, 6)), Mlist, Glist, Slist, 0.1, 2) + assert th.shape == (1, 2, 3) + assert dth.shape == (1, 2, 3) diff --git a/packages/PyTorch/tests/test_ch9_11.py b/packages/PyTorch/tests/test_ch9_11.py new file mode 100644 index 0000000..de1ce8f --- /dev/null +++ b/packages/PyTorch/tests/test_ch9_11.py @@ -0,0 +1,95 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import pytorch_mr.core as pmr + + +def test_time_scaling_examples(): + assert abs(float(pmr.CubicTimeScaling(2, 0.6)) - float(np_mr.CubicTimeScaling(2, 0.6))) < 1e-12 + assert abs(float(pmr.QuinticTimeScaling(2, 0.6)) - float(np_mr.QuinticTimeScaling(2, 0.6))) < 1e-12 + + +def test_joint_trajectory_example_endpoints(): + start = torch.tensor([1.0, 0.0, 0.0, 1.0, 1.0, 0.2, 0.0, 1.0]) + end = torch.tensor([1.2, 0.5, 0.6, 1.1, 2.0, 2.0, 0.9, 1.0]) + traj = pmr.JointTrajectory(start, end, 4, 6, 3).squeeze(0) + assert torch.allclose(traj[0], start, atol=1e-6) + assert torch.allclose(traj[-1], end, atol=1e-6) + traj_np = np_mr.JointTrajectory(start.cpu().numpy(), end.cpu().numpy(), 4, 6, 3) + assert np.allclose(traj.cpu().numpy(), traj_np, atol=1e-6) + + +def test_screw_and_cartesian_trajectory_shapes(): + Xstart = torch.tensor([[1.0, 0.0, 0.0, 1.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 1.0], + [0.0, 0.0, 0.0, 1.0]]) + Xend = torch.tensor([[0.0, 0.0, 1.0, 0.1], + [1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 4.1], + [0.0, 0.0, 0.0, 1.0]]) + screw = pmr.ScrewTrajectory(Xstart, Xend, 5, 4, 3) + cart = pmr.CartesianTrajectory(Xstart, Xend, 5, 4, 5) + assert screw.shape == (1, 4, 4, 4) + assert cart.shape == (1, 4, 4, 4) + assert torch.allclose(screw[:, 0], Xstart.unsqueeze(0), atol=1e-6) + assert torch.allclose(screw[:, -1], Xend.unsqueeze(0), atol=1e-5) + screw_np = np_mr.ScrewTrajectory(Xstart.cpu().numpy(), Xend.cpu().numpy(), 5, 4, 3) + cart_np = np_mr.CartesianTrajectory(Xstart.cpu().numpy(), Xend.cpu().numpy(), 5, 4, 5) + assert np.allclose(screw.squeeze(0).cpu().numpy(), np.stack(screw_np, axis=0), atol=1e-5) + assert np.allclose(cart.squeeze(0).cpu().numpy(), np.stack(cart_np, axis=0), atol=1e-5) + + +def test_computed_torque_and_simulate_control_smoke(): + M01 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.089159], + [0.0, 0.0, 0.0, 1.0]]) + M12 = torch.tensor([[0.0, 0.0, 1.0, 0.28], + [0.0, 1.0, 0.0, 0.13585], + [-1.0, 0.0, 0.0, 0.0], + [0.0, 0.0, 0.0, 1.0]]) + M23 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.1197], + [0.0, 0.0, 1.0, 0.395], + [0.0, 0.0, 0.0, 1.0]]) + M34 = torch.tensor([[1.0, 0.0, 0.0, 0.0], + [0.0, 1.0, 0.0, 0.0], + [0.0, 0.0, 1.0, 0.14225], + [0.0, 0.0, 0.0, 1.0]]) + G1 = torch.diag(torch.tensor([0.010267, 0.010267, 0.00666, 3.7, 3.7, 3.7])) + G2 = torch.diag(torch.tensor([0.22689, 0.22689, 0.0151074, 8.393, 8.393, 8.393])) + G3 = torch.diag(torch.tensor([0.0494433, 0.0494433, 0.004095, 2.275, 2.275, 2.275])) + Mlist = torch.stack([M01, M12, M23, M34], dim=0) + Glist = torch.stack([G1, G2, G3], dim=0) + Slist = torch.tensor([[1.0, 0.0, 1.0, 0.0, 1.0, 0.0], + [0.0, 1.0, 0.0, -0.089, 0.0, 0.0], + [0.0, 1.0, 0.0, -0.089, 0.0, 0.425]]).T + thetalist = torch.tensor([0.1, 0.1, 0.1]) + dthetalist = torch.tensor([0.1, 0.2, 0.3]) + eint = torch.tensor([0.2, 0.2, 0.2]) + g = torch.tensor([0.0, 0.0, -9.8]) + tau = pmr.ComputedTorque( + thetalist=thetalist, dthetalist=dthetalist, eint=eint, g=g, Mlist=Mlist, Glist=Glist, Slist=Slist, + thetalistd=torch.tensor([1.0, 1.0, 1.0]), dthetalistd=torch.tensor([2.0, 1.2, 2.0]), + ddthetalistd=torch.tensor([0.1, 0.1, 0.1]), Kp=1.3, Ki=1.2, Kd=1.1 + ).squeeze(0) + expected_tau = np_mr.ComputedTorque( + thetalist.cpu().numpy(), dthetalist.cpu().numpy(), eint.cpu().numpy(), g.cpu().numpy(), + Mlist.cpu().numpy(), Glist.cpu().numpy(), Slist.cpu().numpy(), + np.array([1.0, 1.0, 1.0]), np.array([2.0, 1.2, 2.0]), np.array([0.1, 0.1, 0.1]), + 1.3, 1.2, 1.1 + ) + assert np.allclose(tau.cpu().numpy(), expected_tau, atol=1e-2) + + traj = pmr.JointTrajectory(thetalist, torch.tensor([1.57, 3.14, 4.71]), 0.1, 5, 5).squeeze(0) + dtraj = torch.zeros_like(traj) + ddtraj = torch.zeros_like(traj) + taumat, thetamat = pmr.SimulateControl( + thetalist=thetalist, dthetalist=dthetalist, g=g, Ftipmat=torch.ones((5, 6)), + Mlist=Mlist, Glist=Glist, Slist=Slist, thetamatd=traj, dthetamatd=dtraj, ddthetamatd=ddtraj, + gtilde=g, Mtildelist=Mlist, Gtildelist=Glist, Kp=20, Ki=10, Kd=18, dt=0.02, intRes=2 + ) + assert taumat.shape == (1, 5, 3) + assert thetamat.shape == (1, 5, 3) diff --git a/packages/Python/setup.py b/packages/Python/setup.py index 5ba8aef..eefb6c9 100644 --- a/packages/Python/setup.py +++ b/packages/Python/setup.py @@ -17,7 +17,7 @@ setup( name = "modern_robotics", version = __version__, - author = "Huan Weng, Mikhail Todes, Jarvis Schultz, Bill Hunt", + author = "Huan Weng, Mikhail Todes, Jarvis Schultz, Bill Hunt, Ayush Gaggar", author_email = "huanweng@u.northwestern.edu", description = ("Modern Robotics: Mechanics, Planning, and Control: Code Library"), license = "MIT", @@ -41,4 +41,4 @@ 'numpy', ], platforms='Linux, Mac, Windows', -) +) \ No newline at end of file