From e138c7b7035025e154f08cfa3b7465bcf8d5f6aa Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Wed, 16 Jul 2025 15:54:49 -0500 Subject: [PATCH 1/7] adding mr_pytorch package --- README.md | 1 + packages/Python/mr_pytorch/__init__.py | 3 + packages/Python/mr_pytorch/__version__.py | 2 + packages/Python/mr_pytorch/core.py | 495 ++++++++++++++++++++++ packages/Python/setup.py | 29 ++ 5 files changed, 530 insertions(+) create mode 100644 packages/Python/mr_pytorch/__init__.py create mode 100644 packages/Python/mr_pytorch/__version__.py create mode 100644 packages/Python/mr_pytorch/core.py diff --git a/README.md b/README.md index f2ad38c..b105f6d 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) 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/Python/mr_pytorch/__init__.py b/packages/Python/mr_pytorch/__init__.py new file mode 100644 index 0000000..aada3de --- /dev/null +++ b/packages/Python/mr_pytorch/__init__.py @@ -0,0 +1,3 @@ +from .__version__ import __version__ + +from .core import * diff --git a/packages/Python/mr_pytorch/__version__.py b/packages/Python/mr_pytorch/__version__.py new file mode 100644 index 0000000..d930acc --- /dev/null +++ b/packages/Python/mr_pytorch/__version__.py @@ -0,0 +1,2 @@ + +__version__ = '1.1.0' diff --git a/packages/Python/mr_pytorch/core.py b/packages/Python/mr_pytorch/core.py new file mode 100644 index 0000000..119b832 --- /dev/null +++ b/packages/Python/mr_pytorch/core.py @@ -0,0 +1,495 @@ +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. +''' + +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) + """ + 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,) + """ + if expc3.ndim == 2: + expc3 = expc3.unsqueeze(-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]]) + """ + 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) + """ + #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.tensor([0, 0, 0, 1], device=V.device, dtype=V.dtype) + bottom_row = bottom_row.expand(V.size(0), 1, 4) # (N, 1, 4) + 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) + """ + return NotImplementedError("ScrewToAxis is not yet tested.") + if q.ndim == 2: + q = q.unsqueeze(-1) + if s.ndim == 2: + s = s.unsqueeze(-1) + if h.ndim == 1: + h = h.unsqueeze(-1) + + return torch.cat([s, torch.cross(q.squeeze(-1), s.squeeze(-1), dim=-1) + h * s.squeeze(-1)], dim=-1).unsqueeze(-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,) + """ + return NotImplementedError("AxisAng6 is not yet tested.") + theta = torch.linalg.norm(expc6[:, :3], dim=-1, keepdim=True) + mask = NearZero(theta) + theta[mask] = torch.linalg.norm(expc6[:, 3:], dim=-1, keepdim=True) + return (expc6 / theta).squeeze(-1), theta.squeeze(-1) + +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) + """ + 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) + """ + return NotImplementedError("ProjectToSO3 is not yet implemented.") + +def ProjectToSE3(T: torch.Tensor): + return NotImplementedError("ProjectToSE3 is not yet implemented.") + +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,) + """ + return NotImplementedError("DistanceToSO3 is not yet implemented.") + +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,) + """ + return NotImplementedError("DistanceToSE3 is not yet implemented.") + +def TestIfSO3(R: torch.Tensor): + return NotImplementedError("TestIfSO3 is not yet implemented.") + +def TestIfSE3(T: torch.Tensor): + return NotImplementedError("TestIfSE3 is not yet implemented.") + +def FKinBody(M: torch.Tensor, Blist: torch.Tensor, thetalist: torch.Tensor): + return NotImplementedError("FKinBody is not yet implemented.") + +def FKinSpace(M: torch.Tensor, Slist: torch.Tensor, thetalist: torch.Tensor): + return NotImplementedError("FKinSpace is not yet implemented.") + +def JacobianBody(Blist: torch.Tensor, thetalist: torch.Tensor): + return NotImplementedError("JacobianBody is not yet implemented.") + +def JacobianSpace(Slist: torch.Tensor, thetalist: torch.Tensor): + return NotImplementedError("JacobianSpace is not yet implemented.") + +def IKinBody(Blist: torch.Tensor, M: torch.Tensor, T: torch.Tensor, thetalist0: torch.Tensor, eomg=1e-6, ev=1e-6): + return NotImplementedError("IKinBody is not yet implemented.") + +def IKinSpace(Slist: torch.Tensor, M: torch.Tensor, T: torch.Tensor, thetalist0: torch.Tensor, eomg=1e-6, ev=1e-6): + return NotImplementedError("IKinSpace is not yet implemented.") + +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) + """ + return NotImplementedError("ad is not yet tested.") + if V.ndim == 1: + V = V.unsqueeze(0) + if V.ndim == 2: + V = V.unsqueeze(-1) + omg = V[:, :3] + v = V[:, 3:] + so3mat = VecToso3(omg) + return torch.cat([torch.cat([so3mat, v], dim=-1), + torch.cat([VecToso3(v), so3mat], 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): + return NotImplementedError("InverseDynamics is not yet implemented.") + +def MassMatrix(thetalist: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + return NotImplementedError("MassMatrix is not yet implemented.") + +def VelQuadraticForces(thetalist: torch.Tensor, dthetalist: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + return NotImplementedError("VelQuadraticForces is not yet implemented.") + +def GravityForces(thetalist: torch.Tensor, g: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + return NotImplementedError("GravityForces is not yet implemented.") + +def EndEffectorForces(thetalist: torch.Tensor, Ftip: torch.Tensor, Mlist: torch.Tensor, Glist: torch.Tensor, Slist: torch.Tensor): + return NotImplementedError("EndEffectorForces is not yet implemented.") + +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): + return NotImplementedError("ForwardDynamics is not yet implemented.") + +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): + return NotImplementedError("InverseDynamicsTrajectory is not yet implemented.") + +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): + return NotImplementedError("ForwardDynamicsTrajectory is not yet implemented.") + +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): + return NotImplementedError("JointTrajectory is not yet implemented.") + +def ScrewTrajectory(Xstart: torch.Tensor, Xend: torch.Tensor, Tf: float, N: int, method): + return NotImplementedError("ScrewTrajectory is not yet implemented.") + +def CartesianTrajectory(Xstart: torch.Tensor, Xend: torch.Tensor, Tf: float, N: int, method): + return NotImplementedError("CartesianTrajectory is not yet implemented.") + +def ComputedTorque(**kwargs): + return NotImplementedError("ComputedTorque is not yet implemented.") + +def SimulateControl(**kwargs): + return NotImplementedError("SimulateControl is not yet implemented.") diff --git a/packages/Python/setup.py b/packages/Python/setup.py index 5ba8aef..a14fc3e 100644 --- a/packages/Python/setup.py +++ b/packages/Python/setup.py @@ -42,3 +42,32 @@ ], platforms='Linux, Mac, Windows', ) + +setup( + name="mr_pytorch", + version=__version__, + author = "Ayush Gaggar", + author_email = "agaggar@u.northwestern.edu", + description = ("PyTorch version of Modern Robotics Code, useful for batched, learning, and differentiable applications"), + license = "MIT", + long_description = long_description, + long_description_content_type='text/markdown', + keywords = "kinematics robotics dynamics", + url = "http://modernrobotics.org/", + packages=['mr_pytorch'], + classifiers=[ + "Development Status :: 3 - Alpha", + "Intended Audience :: Education", + "Intended Audience :: Science/Research", + "License :: OSI Approved :: MIT License", + "Natural Language :: English", + "Programming Language :: Python :: 2", + "Programming Language :: Python :: 3", + "Topic :: Education", + "Topic :: Scientific/Engineering", + ], + install_requires=[ + 'numpy', + ], + platforms='Linux, Mac, Windows', +) From 8508a156c73477433fbef10e5ce2e17e86e7e494 Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Wed, 16 Jul 2025 16:02:04 -0500 Subject: [PATCH 2/7] updated setup.py file --- .../pytorch_mr}/__init__.py | 0 .../pytorch_mr}/__version__.py | 0 .../pytorch_mr}/core.py | 0 packages/Python/setup.py | 33 ++----------------- 4 files changed, 2 insertions(+), 31 deletions(-) rename packages/Python/{mr_pytorch => modern_robotics/pytorch_mr}/__init__.py (100%) rename packages/Python/{mr_pytorch => modern_robotics/pytorch_mr}/__version__.py (100%) rename packages/Python/{mr_pytorch => modern_robotics/pytorch_mr}/core.py (100%) diff --git a/packages/Python/mr_pytorch/__init__.py b/packages/Python/modern_robotics/pytorch_mr/__init__.py similarity index 100% rename from packages/Python/mr_pytorch/__init__.py rename to packages/Python/modern_robotics/pytorch_mr/__init__.py diff --git a/packages/Python/mr_pytorch/__version__.py b/packages/Python/modern_robotics/pytorch_mr/__version__.py similarity index 100% rename from packages/Python/mr_pytorch/__version__.py rename to packages/Python/modern_robotics/pytorch_mr/__version__.py diff --git a/packages/Python/mr_pytorch/core.py b/packages/Python/modern_robotics/pytorch_mr/core.py similarity index 100% rename from packages/Python/mr_pytorch/core.py rename to packages/Python/modern_robotics/pytorch_mr/core.py diff --git a/packages/Python/setup.py b/packages/Python/setup.py index a14fc3e..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,33 +41,4 @@ 'numpy', ], platforms='Linux, Mac, Windows', -) - -setup( - name="mr_pytorch", - version=__version__, - author = "Ayush Gaggar", - author_email = "agaggar@u.northwestern.edu", - description = ("PyTorch version of Modern Robotics Code, useful for batched, learning, and differentiable applications"), - license = "MIT", - long_description = long_description, - long_description_content_type='text/markdown', - keywords = "kinematics robotics dynamics", - url = "http://modernrobotics.org/", - packages=['mr_pytorch'], - classifiers=[ - "Development Status :: 3 - Alpha", - "Intended Audience :: Education", - "Intended Audience :: Science/Research", - "License :: OSI Approved :: MIT License", - "Natural Language :: English", - "Programming Language :: Python :: 2", - "Programming Language :: Python :: 3", - "Topic :: Education", - "Topic :: Scientific/Engineering", - ], - install_requires=[ - 'numpy', - ], - platforms='Linux, Mac, Windows', -) +) \ No newline at end of file From 202bc07b92fa5c0ee3164989a2daa08e7ab44c08 Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Wed, 16 Jul 2025 16:02:42 -0500 Subject: [PATCH 3/7] specified import --- README.md | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/README.md b/README.md index b105f6d..74ce0cf 100644 --- a/README.md +++ b/README.md @@ -20,7 +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) +* [PyTorch version](https://github.com/Agaggar/PyTorch_MR), to be used as `from modern_robotics 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. From 2052a58f3c478f5b8720d3c7a70ba44880e241d1 Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Tue, 31 Mar 2026 17:09:07 -0500 Subject: [PATCH 4/7] complete implementation of numpy package --- .../Python/modern_robotics/pytorch_mr/core.py | 630 ++++++++++++++++-- 1 file changed, 579 insertions(+), 51 deletions(-) diff --git a/packages/Python/modern_robotics/pytorch_mr/core.py b/packages/Python/modern_robotics/pytorch_mr/core.py index 119b832..6025f50 100644 --- a/packages/Python/modern_robotics/pytorch_mr/core.py +++ b/packages/Python/modern_robotics/pytorch_mr/core.py @@ -15,8 +15,34 @@ 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 @@ -68,6 +94,8 @@ def so3ToVec(so3mat: torch.Tensor): :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], @@ -83,8 +111,9 @@ def AxisAng3(expc3: torch.Tensor): :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) + 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) @@ -106,6 +135,7 @@ def MatrixExp3(so3mat: torch.Tensor): [-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) @@ -122,6 +152,7 @@ def MatrixLog3(R: torch.Tensor): :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) @@ -204,8 +235,7 @@ def VecTose3(V: torch.Tensor): 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.tensor([0, 0, 0, 1], device=V.device, dtype=V.dtype) - bottom_row = bottom_row.expand(V.size(0), 1, 4) # (N, 1, 4) + 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) @@ -251,15 +281,15 @@ def ScrewToAxis(q: torch.Tensor, s: torch.Tensor, h: torch.Tensor): :param h: The pitch of the screw, shape (N,) :return: The axis representation of the screw, shape (N, 6) """ - return NotImplementedError("ScrewToAxis is not yet tested.") - if q.ndim == 2: - q = q.unsqueeze(-1) - if s.ndim == 2: - s = s.unsqueeze(-1) - if h.ndim == 1: - h = h.unsqueeze(-1) - - return torch.cat([s, torch.cross(q.squeeze(-1), s.squeeze(-1), dim=-1) + h * s.squeeze(-1)], dim=-1).unsqueeze(-1) + 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 @@ -270,11 +300,12 @@ def AxisAng6(expc6: torch.Tensor): :return S: The corresponding normalized screw axis, shape (N, 6) :return theta: The distance traveled along/about S, shape (N, 1) or (N,) """ - return NotImplementedError("AxisAng6 is not yet tested.") - theta = torch.linalg.norm(expc6[:, :3], dim=-1, keepdim=True) + expc6 = _as_batch_vec(expc6, 6) + theta = torch.linalg.norm(expc6[:, :3], dim=-1) mask = NearZero(theta) - theta[mask] = torch.linalg.norm(expc6[:, 3:], dim=-1, keepdim=True) - return (expc6 / theta).squeeze(-1), theta.squeeze(-1) + 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 @@ -311,6 +342,7 @@ def MatrixLog6(T: torch.Tensor): :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: @@ -346,10 +378,30 @@ def ProjectToSO3(R: torch.Tensor): :param R: A rotation matrix, shape (N, 3, 3) :return: The projected rotation matrix, shape (N, 3, 3) """ - return NotImplementedError("ProjectToSO3 is not yet implemented.") + 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): - return NotImplementedError("ProjectToSE3 is not yet implemented.") + """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 @@ -358,7 +410,15 @@ def DistanceToSO3(R: torch.Tensor): :param R: A matrix, shape (N, 3, 3) :return: A quantity describing the distance of mat from the SO(3) manifold, shape (N,) """ - return NotImplementedError("DistanceToSO3 is not yet implemented.") + 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 @@ -366,32 +426,190 @@ def DistanceToSE3(T: torch.Tensor): :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 """ - return NotImplementedError("DistanceToSE3 is not yet implemented.") + 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): - return NotImplementedError("TestIfSO3 is not yet implemented.") + """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): - return NotImplementedError("TestIfSE3 is not yet implemented.") + """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): - return NotImplementedError("FKinBody is not yet implemented.") + """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): - return NotImplementedError("FKinSpace is not yet implemented.") + """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): - return NotImplementedError("JacobianBody is not yet implemented.") + """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): - return NotImplementedError("JacobianSpace is not yet implemented.") + """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): - return NotImplementedError("IKinBody is not yet implemented.") + """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): - return NotImplementedError("IKinSpace is not yet implemented.") + """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 @@ -399,34 +617,144 @@ def ad(V: torch.Tensor): :param V: A spatial velocity vector, shape (N, 6) :return: The adjoint representation of V, shape (N, 6, 6) """ - return NotImplementedError("ad is not yet tested.") - if V.ndim == 1: - V = V.unsqueeze(0) - if V.ndim == 2: - V = V.unsqueeze(-1) - omg = V[:, :3] - v = V[:, 3:] - so3mat = VecToso3(omg) - return torch.cat([torch.cat([so3mat, v], dim=-1), - torch.cat([VecToso3(v), so3mat], dim=-1)], dim=-2) + 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): - return NotImplementedError("InverseDynamics is not yet implemented.") + """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): - return NotImplementedError("MassMatrix is not yet implemented.") + """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): - return NotImplementedError("VelQuadraticForces is not yet implemented.") + """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): - return NotImplementedError("GravityForces is not yet implemented.") + """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): - return NotImplementedError("EndEffectorForces is not yet implemented.") + """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): - return NotImplementedError("ForwardDynamics is not yet implemented.") + """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. @@ -440,10 +768,72 @@ def EulerStep(thetalist: torch.Tensor, dthetalist: torch.Tensor, ddthetalist: to 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): - return NotImplementedError("InverseDynamicsTrajectory is not yet implemented.") + """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): - return NotImplementedError("ForwardDynamicsTrajectory is not yet implemented.") + """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 @@ -480,16 +870,154 @@ def QuinticTimeScaling(Tf, t): + 6 * (1.0 * t / Tf) ** 5 def JointTrajectory(thetalist: torch.Tensor, thetaend: torch.Tensor, Tf: float, N: int, method): - return NotImplementedError("JointTrajectory is not yet implemented.") + """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): - return NotImplementedError("ScrewTrajectory is not yet implemented.") + """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): - return NotImplementedError("CartesianTrajectory is not yet implemented.") + """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): - return NotImplementedError("ComputedTorque is not yet implemented.") + """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): - return NotImplementedError("SimulateControl is not yet implemented.") + """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 From dcc55ae57bacb2d7e26c0d8a12420344dbd74950 Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Tue, 31 Mar 2026 17:09:24 -0500 Subject: [PATCH 5/7] added unit tests WRITTEN BY AI --- .../pytorch_mr/tests/README.md | 40 +++++ .../pytorch_mr/tests/conftest.py | 128 +++++++++++++ .../pytorch_mr/tests/run_tests_by_chapter.py | 77 ++++++++ .../tests/test_ch1_3_numpy_parity.py | 170 ++++++++++++++++++ .../pytorch_mr/tests/test_ch3.py | 67 +++++++ .../pytorch_mr/tests/test_ch4_6.py | 79 ++++++++ .../pytorch_mr/tests/test_ch8.py | 105 +++++++++++ .../pytorch_mr/tests/test_ch9_11.py | 95 ++++++++++ 8 files changed, 761 insertions(+) create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/README.md create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/conftest.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py create mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/README.md b/packages/Python/modern_robotics/pytorch_mr/tests/README.md new file mode 100644 index 0000000..535d0d6 --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/README.md @@ -0,0 +1,40 @@ +WARNING: All unit tests have been written and checked by AI. +To run unit tests, run the following in the directory `./packages/Python/`: +`PYTHONPATH=. python modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py`. +Expected output: +..................... [100%] +21 passed in 0.05s +...... [100%] +6 passed in 0.03s +.... [100%] +4 passed in 0.05s +... [100%] +3 passed in 0.09s +.... [100%] +4 passed in 0.20s + +=== Chapter test summary === + ch1_3 | PASS | collected= 21 | modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py + ch3 | PASS | collected= 6 | modern_robotics/pytorch_mr/tests/test_ch3.py + ch4_6 | PASS | collected= 4 | modern_robotics/pytorch_mr/tests/test_ch4_6.py + ch8 | PASS | collected= 3 | modern_robotics/pytorch_mr/tests/test_ch8.py +ch9_11 | PASS | collected= 4 | modern_robotics/pytorch_mr/tests/test_ch9_11.py +TOTAL | PASS | collected= 38 + +Specifically, Cursor Pro was used with the Plan phase and the Agent set to Auto (as of March 2026). The exact prompt to generate these tests were: + +Write a set of unit tests for each pytorch function. At the bare minimum, the tests should include and pass the examples mentioned in the comments of the numpy version. Make sure to test for possible edge cases (such as angle wrapping, near zero issues, etc.) Please put all unit tests in a separate directory under pytorch_mr. + +If you are more than 90% confident on the next steps to implement after reviewing the specified plan, proceed with writing code. Feel free to make design decisions along the way (such as improvements to speed, broadcasting, etc.) but make sure to explicitly comment on this. Otherwise, ask some follow up questions to gain a better understanding on the course of action. If implementing any design decisions, be sure to comment decisions in a careful manner. + +Based on your current understanding of the codebase, the goal of the project, and the feature/changes being described, suggest a course of action that would best accomplish this goal. +Act as another researcher with deep technical expertise that pokes holes in the overall methodology, motivation, and implementation. There's no need to "over-engineer" the approach, but there is a definite need to make this approach fully thought through. +There might not be a need to begin programming changes right away, but if you and I are both more than 90% confident with the decision and implementation plan, proceed with the code. If you do start coding, give me a summary of changes as well as some terminal commands to test the performance of the changes. + + +After initially creating the unit tests, the follow up prompt given was: +Also do a second pass on unit tests that enforces stricter shape-normalization with parametrized tests. Also add any other tests that you think might be worth checking. + +For all unit tests, check that the output between the function implemented in pytorch matches the function call to the numpy version of modern robotics. + +Finally, write a script that calls the unit tests and prints an output detailing which tests were run from which chapter, and how many tests passed or failed. \ No newline at end of file diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/conftest.py b/packages/Python/modern_robotics/pytorch_mr/tests/conftest.py new file mode 100644 index 0000000..ff76971 --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/conftest.py @@ -0,0 +1,128 @@ +import os +import sys +from typing import Any, Callable, Sequence + +import numpy as np +import pytest +import torch + + +# Ensure local `packages/Python` is importable (not site-packages). +_HERE = os.path.dirname(__file__) +_PKG_PYTHON_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) +if _PKG_PYTHON_ROOT not in sys.path: + sys.path.insert(0, _PKG_PYTHON_ROOT) + + +import modern_robotics.core as mr_np # noqa: E402 +import modern_robotics.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. + """ + # Determine batch size from first tensor argument (if any) + 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/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py b/packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py new file mode 100644 index 0000000..88e538a --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py @@ -0,0 +1,77 @@ +#!/usr/bin/env python3 +import os +import re +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): + p = subprocess.run( + [sys.executable, "-m", "pytest", "--collect-only", "-q", *pytest_args], + stdout=subprocess.PIPE, + stderr=subprocess.STDOUT, + text=True, + env=env, + ) + # Each collected test item prints one line like: path::test_name + 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): + p = subprocess.run( + [sys.executable, "-m", "pytest", "-q", *pytest_args], + env=env, + ) + return p.returncode + + +def main(): + here = os.path.abspath(os.path.dirname(__file__)) + # packages/Python (so `modern_robotics/...` paths exist) + pkg_root = os.path.abspath(os.path.join(here, "..", "..", "..")) + env = dict(os.environ) + env["PYTHONPATH"] = pkg_root + (os.pathsep + 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, pkg_root) + collected = _collect_count([rel], env) + exit_code = _run([rel], env) + 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/Python/modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py new file mode 100644 index 0000000..c1f7334 --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/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/Python/modern_robotics/pytorch_mr/tests/test_ch3.py b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py new file mode 100644 index 0000000..eec4508 --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py @@ -0,0 +1,67 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import modern_robotics.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/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py new file mode 100644 index 0000000..5685a5f --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py @@ -0,0 +1,79 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import modern_robotics.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/Python/modern_robotics/pytorch_mr/tests/test_ch8.py b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py new file mode 100644 index 0000000..2c01cea --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py @@ -0,0 +1,105 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import modern_robotics.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/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py new file mode 100644 index 0000000..9735a44 --- /dev/null +++ b/packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py @@ -0,0 +1,95 @@ +import numpy as np +import torch + +import modern_robotics.core as np_mr +import modern_robotics.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) From d9067b336e365f4121cbe50cdf02fe6dcc6e9bce Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Tue, 31 Mar 2026 17:20:54 -0500 Subject: [PATCH 6/7] changed directory structure to have a separate pytorch folder --- README.md | 2 +- packages/PyTorch/README.md | 50 +++++++++++++++++++ packages/PyTorch/pyproject.toml | 3 ++ .../pytorch_mr/__init__.py | 0 .../pytorch_mr/__version__.py | 0 .../pytorch_mr/core.py | 0 packages/PyTorch/pytorch_mr/tests/README.md | 25 ++++++++++ packages/PyTorch/pytorch_mr/tests/__init__.py | 1 + .../pytorch_mr/tests/conftest.py | 24 +++++---- .../pytorch_mr/tests/run_tests_by_chapter.py | 15 +++--- .../tests/test_ch1_3_numpy_parity.py | 0 .../pytorch_mr/tests/test_ch3.py | 2 +- .../pytorch_mr/tests/test_ch4_6.py | 2 +- .../pytorch_mr/tests/test_ch8.py | 2 +- .../pytorch_mr/tests/test_ch9_11.py | 2 +- packages/PyTorch/setup.py | 47 +++++++++++++++++ .../pytorch_mr/tests/README.md | 40 --------------- 17 files changed, 154 insertions(+), 61 deletions(-) create mode 100644 packages/PyTorch/README.md create mode 100644 packages/PyTorch/pyproject.toml rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/__init__.py (100%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/__version__.py (100%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/core.py (100%) create mode 100644 packages/PyTorch/pytorch_mr/tests/README.md create mode 100644 packages/PyTorch/pytorch_mr/tests/__init__.py rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/conftest.py (81%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/run_tests_by_chapter.py (81%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/test_ch1_3_numpy_parity.py (100%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/test_ch3.py (98%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/test_ch4_6.py (98%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/test_ch8.py (99%) rename packages/{Python/modern_robotics => PyTorch}/pytorch_mr/tests/test_ch9_11.py (99%) create mode 100644 packages/PyTorch/setup.py delete mode 100644 packages/Python/modern_robotics/pytorch_mr/tests/README.md diff --git a/README.md b/README.md index 74ce0cf..7c6fcc8 100644 --- a/README.md +++ b/README.md @@ -20,7 +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), to be used as `from modern_robotics import pytorch_mr as pmr` +* [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/README.md b/packages/PyTorch/README.md new file mode 100644 index 0000000..7bbbd21 --- /dev/null +++ b/packages/PyTorch/README.md @@ -0,0 +1,50 @@ +# 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 + +From this directory (`packages/PyTorch`): + +```bash +pip install -e . +``` + +Dependencies: `numpy`, `torch` (see `setup.py`). + +## Import + +```python +import pytorch_mr as pmr +# or +from pytorch_mr.core import MatrixExp3, FKinBody +``` + +## Tests (NumPy parity) + +Tests compare against the reference `modern_robotics` package. Install the NumPy library from the sibling directory or from PyPI: + +```bash +pip install -e ../Python # local reference, from packages/PyTorch +# or +pip install modern_robotics +pip install pytest +``` + +Run tests with both package roots on `PYTHONPATH` (from `packages/` in this repo): + +```bash +cd packages +PYTHONPATH=PyTorch:Python python -m pytest PyTorch/pytorch_mr/tests -q +``` + +Or use the chapter summary script: + +```bash +cd packages +PYTHONPATH=PyTorch:Python python PyTorch/pytorch_mr/tests/run_tests_by_chapter.py +``` + +## Publishing to PyPI (optional) + +Build and upload with standard `twine` / `build` workflows; ensure the package name `pytorch_mr` is available on your target index. diff --git a/packages/PyTorch/pyproject.toml b/packages/PyTorch/pyproject.toml new file mode 100644 index 0000000..4a85092 --- /dev/null +++ b/packages/PyTorch/pyproject.toml @@ -0,0 +1,3 @@ +[build-system] +requires = ["setuptools>=61", "wheel"] +build-backend = "setuptools.build_meta" diff --git a/packages/Python/modern_robotics/pytorch_mr/__init__.py b/packages/PyTorch/pytorch_mr/__init__.py similarity index 100% rename from packages/Python/modern_robotics/pytorch_mr/__init__.py rename to packages/PyTorch/pytorch_mr/__init__.py diff --git a/packages/Python/modern_robotics/pytorch_mr/__version__.py b/packages/PyTorch/pytorch_mr/__version__.py similarity index 100% rename from packages/Python/modern_robotics/pytorch_mr/__version__.py rename to packages/PyTorch/pytorch_mr/__version__.py diff --git a/packages/Python/modern_robotics/pytorch_mr/core.py b/packages/PyTorch/pytorch_mr/core.py similarity index 100% rename from packages/Python/modern_robotics/pytorch_mr/core.py rename to packages/PyTorch/pytorch_mr/core.py diff --git a/packages/PyTorch/pytorch_mr/tests/README.md b/packages/PyTorch/pytorch_mr/tests/README.md new file mode 100644 index 0000000..cd172a4 --- /dev/null +++ b/packages/PyTorch/pytorch_mr/tests/README.md @@ -0,0 +1,25 @@ +# Unit tests for `pytorch_mr` + +Tests live next to the PyTorch package under `packages/PyTorch/pytorch_mr/tests/`. + +They compare outputs to the NumPy reference `modern_robotics` (`packages/Python`). Ensure both are importable: + +- From repo `packages/` directory: + +```bash +PYTHONPATH=PyTorch:Python python -m pytest PyTorch/pytorch_mr/tests -q +``` + +- Chapter summary: + +```bash +PYTHONPATH=PyTorch:Python python PyTorch/pytorch_mr/tests/run_tests_by_chapter.py +``` + +After `pip install -e .` from `packages/PyTorch`, `pip install pytest`, and a NumPy reference (`pip install modern_robotics` or `pip install -e ../Python`), run: + +```bash +PYTHONPATH=/path/to/repo/packages/Python python -m pytest --pyargs pytorch_mr.tests -q +``` + +(Use `PYTHONPATH` to point at the local `modern_robotics` tree when you are not using the PyPI package.) diff --git a/packages/PyTorch/pytorch_mr/tests/__init__.py b/packages/PyTorch/pytorch_mr/tests/__init__.py new file mode 100644 index 0000000..c919acf --- /dev/null +++ b/packages/PyTorch/pytorch_mr/tests/__init__.py @@ -0,0 +1 @@ +# Test package for pytorch_mr (optional discovery via pytest --pyargs pytorch_mr.tests) diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/conftest.py b/packages/PyTorch/pytorch_mr/tests/conftest.py similarity index 81% rename from packages/Python/modern_robotics/pytorch_mr/tests/conftest.py rename to packages/PyTorch/pytorch_mr/tests/conftest.py index ff76971..03e4c39 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/conftest.py +++ b/packages/PyTorch/pytorch_mr/tests/conftest.py @@ -6,16 +6,24 @@ import pytest import torch - -# Ensure local `packages/Python` is importable (not site-packages). _HERE = os.path.dirname(__file__) -_PKG_PYTHON_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) -if _PKG_PYTHON_ROOT not in sys.path: - sys.path.insert(0, _PKG_PYTHON_ROOT) +# Prefer installed packages; fall back to sibling repo layout (packages/PyTorch + packages/Python). +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 -import modern_robotics.core as mr_np # noqa: E402 -import modern_robotics.pytorch_mr.core as mr_t # noqa: E402 +try: + import pytorch_mr.core as mr_t # noqa: E402 +except ImportError: + _PKG_TORCH_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) + if _PKG_TORCH_ROOT not in sys.path: + sys.path.insert(0, _PKG_TORCH_ROOT) + import pytorch_mr.core as mr_t # noqa: E402 @pytest.fixture(scope="session") @@ -94,7 +102,6 @@ def apply_numpy_per_batch( - If shape is (N, ...), iterate over axis 0. - If shape is (...), treat as single item. """ - # Determine batch size from first tensor argument (if any) batch = None for a in args_t: if isinstance(a, torch.Tensor) and a.ndim >= 2: @@ -125,4 +132,3 @@ def apply_numpy_per_batch( np_args_i.append(a) outs.append(f_np(*np_args_i)) return outs - diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py b/packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py similarity index 81% rename from packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py rename to packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py index 88e538a..6d8e58c 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py +++ b/packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py @@ -1,6 +1,5 @@ #!/usr/bin/env python3 import os -import re import subprocess import sys from dataclasses import dataclass @@ -22,7 +21,6 @@ def _collect_count(pytest_args, env): text=True, env=env, ) - # Each collected test item prints one line like: path::test_name 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) @@ -38,10 +36,14 @@ def _run(pytest_args, env): def main(): here = os.path.abspath(os.path.dirname(__file__)) - # packages/Python (so `modern_robotics/...` paths exist) - pkg_root = os.path.abspath(os.path.join(here, "..", "..", "..")) + # packages/ — parent of PyTorch/ and Python/ + packages_dir = os.path.abspath(os.path.join(here, "..", "..", "..")) + torch_pkg_root = os.path.join(packages_dir, "PyTorch") + python_pkg_root = os.path.join(packages_dir, "Python") env = dict(os.environ) - env["PYTHONPATH"] = pkg_root + (os.pathsep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "") + sep = os.pathsep + extra = torch_pkg_root + 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")), @@ -56,7 +58,7 @@ def main(): any_failed = False for name, path in chapters: - rel = os.path.relpath(path, pkg_root) + rel = os.path.relpath(path, packages_dir) collected = _collect_count([rel], env) exit_code = _run([rel], env) runs.append(ChapterRun(name=name, path=rel, collected=collected, exit_code=exit_code)) @@ -74,4 +76,3 @@ def main(): if __name__ == "__main__": raise SystemExit(main()) - diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py b/packages/PyTorch/pytorch_mr/tests/test_ch1_3_numpy_parity.py similarity index 100% rename from packages/Python/modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py rename to packages/PyTorch/pytorch_mr/tests/test_ch1_3_numpy_parity.py diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py b/packages/PyTorch/pytorch_mr/tests/test_ch3.py similarity index 98% rename from packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py rename to packages/PyTorch/pytorch_mr/tests/test_ch3.py index eec4508..165a34e 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch3.py +++ b/packages/PyTorch/pytorch_mr/tests/test_ch3.py @@ -2,7 +2,7 @@ import torch import modern_robotics.core as np_mr -import modern_robotics.pytorch_mr.core as pmr +import pytorch_mr.core as pmr def test_screw_to_axis_example(): diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py b/packages/PyTorch/pytorch_mr/tests/test_ch4_6.py similarity index 98% rename from packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py rename to packages/PyTorch/pytorch_mr/tests/test_ch4_6.py index 5685a5f..1f6179a 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch4_6.py +++ b/packages/PyTorch/pytorch_mr/tests/test_ch4_6.py @@ -2,7 +2,7 @@ import torch import modern_robotics.core as np_mr -import modern_robotics.pytorch_mr.core as pmr +import pytorch_mr.core as pmr def test_fkin_body_example(): diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py b/packages/PyTorch/pytorch_mr/tests/test_ch8.py similarity index 99% rename from packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py rename to packages/PyTorch/pytorch_mr/tests/test_ch8.py index 2c01cea..bfe7d4c 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch8.py +++ b/packages/PyTorch/pytorch_mr/tests/test_ch8.py @@ -2,7 +2,7 @@ import torch import modern_robotics.core as np_mr -import modern_robotics.pytorch_mr.core as pmr +import pytorch_mr.core as pmr def robot3(): diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py b/packages/PyTorch/pytorch_mr/tests/test_ch9_11.py similarity index 99% rename from packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py rename to packages/PyTorch/pytorch_mr/tests/test_ch9_11.py index 9735a44..de1ce8f 100644 --- a/packages/Python/modern_robotics/pytorch_mr/tests/test_ch9_11.py +++ b/packages/PyTorch/pytorch_mr/tests/test_ch9_11.py @@ -2,7 +2,7 @@ import torch import modern_robotics.core as np_mr -import modern_robotics.pytorch_mr.core as pmr +import pytorch_mr.core as pmr def test_time_scaling_examples(): diff --git a/packages/PyTorch/setup.py b/packages/PyTorch/setup.py new file mode 100644 index 0000000..7da9231 --- /dev/null +++ b/packages/PyTorch/setup.py @@ -0,0 +1,47 @@ +from setuptools import find_packages, setup + +exec(open('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="Huan Weng, Mikhail Todes, Jarvis Schultz, Bill Hunt, Ayush Gaggar", + author_email="huanweng@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", +) diff --git a/packages/Python/modern_robotics/pytorch_mr/tests/README.md b/packages/Python/modern_robotics/pytorch_mr/tests/README.md deleted file mode 100644 index 535d0d6..0000000 --- a/packages/Python/modern_robotics/pytorch_mr/tests/README.md +++ /dev/null @@ -1,40 +0,0 @@ -WARNING: All unit tests have been written and checked by AI. -To run unit tests, run the following in the directory `./packages/Python/`: -`PYTHONPATH=. python modern_robotics/pytorch_mr/tests/run_tests_by_chapter.py`. -Expected output: -..................... [100%] -21 passed in 0.05s -...... [100%] -6 passed in 0.03s -.... [100%] -4 passed in 0.05s -... [100%] -3 passed in 0.09s -.... [100%] -4 passed in 0.20s - -=== Chapter test summary === - ch1_3 | PASS | collected= 21 | modern_robotics/pytorch_mr/tests/test_ch1_3_numpy_parity.py - ch3 | PASS | collected= 6 | modern_robotics/pytorch_mr/tests/test_ch3.py - ch4_6 | PASS | collected= 4 | modern_robotics/pytorch_mr/tests/test_ch4_6.py - ch8 | PASS | collected= 3 | modern_robotics/pytorch_mr/tests/test_ch8.py -ch9_11 | PASS | collected= 4 | modern_robotics/pytorch_mr/tests/test_ch9_11.py -TOTAL | PASS | collected= 38 - -Specifically, Cursor Pro was used with the Plan phase and the Agent set to Auto (as of March 2026). The exact prompt to generate these tests were: - -Write a set of unit tests for each pytorch function. At the bare minimum, the tests should include and pass the examples mentioned in the comments of the numpy version. Make sure to test for possible edge cases (such as angle wrapping, near zero issues, etc.) Please put all unit tests in a separate directory under pytorch_mr. - -If you are more than 90% confident on the next steps to implement after reviewing the specified plan, proceed with writing code. Feel free to make design decisions along the way (such as improvements to speed, broadcasting, etc.) but make sure to explicitly comment on this. Otherwise, ask some follow up questions to gain a better understanding on the course of action. If implementing any design decisions, be sure to comment decisions in a careful manner. - -Based on your current understanding of the codebase, the goal of the project, and the feature/changes being described, suggest a course of action that would best accomplish this goal. -Act as another researcher with deep technical expertise that pokes holes in the overall methodology, motivation, and implementation. There's no need to "over-engineer" the approach, but there is a definite need to make this approach fully thought through. -There might not be a need to begin programming changes right away, but if you and I are both more than 90% confident with the decision and implementation plan, proceed with the code. If you do start coding, give me a summary of changes as well as some terminal commands to test the performance of the changes. - - -After initially creating the unit tests, the follow up prompt given was: -Also do a second pass on unit tests that enforces stricter shape-normalization with parametrized tests. Also add any other tests that you think might be worth checking. - -For all unit tests, check that the output between the function implemented in pytorch matches the function call to the numpy version of modern robotics. - -Finally, write a script that calls the unit tests and prints an output detailing which tests were run from which chapter, and how many tests passed or failed. \ No newline at end of file From b9b6cd5ea5c0c4b783e086cbe0a486de26d58396 Mon Sep 17 00:00:00 2001 From: Ayush Gaggar Date: Tue, 31 Mar 2026 18:25:01 -0500 Subject: [PATCH 7/7] changing dir structure and files to upload code to pip --- packages/PyTorch/.gitignore | 8 ++++ packages/PyTorch/LICENSE | 21 ++++++++ packages/PyTorch/README.md | 30 +++++++----- packages/PyTorch/pyproject.toml | 48 +++++++++++++++++++ packages/PyTorch/pytorch_mr/__version__.py | 2 - packages/PyTorch/pytorch_mr/tests/README.md | 25 ---------- packages/PyTorch/pytorch_mr/tests/__init__.py | 1 - packages/PyTorch/setup.py | 9 ++-- .../PyTorch/{ => src}/pytorch_mr/__init__.py | 0 .../PyTorch/src/pytorch_mr/__version__.py | 2 + packages/PyTorch/{ => src}/pytorch_mr/core.py | 0 packages/PyTorch/tests/README.md | 26 ++++++++++ .../{pytorch_mr => }/tests/conftest.py | 10 ++-- .../tests/run_tests_by_chapter.py | 18 +++---- .../tests/test_ch1_3_numpy_parity.py | 0 .../{pytorch_mr => }/tests/test_ch3.py | 0 .../{pytorch_mr => }/tests/test_ch4_6.py | 0 .../{pytorch_mr => }/tests/test_ch8.py | 0 .../{pytorch_mr => }/tests/test_ch9_11.py | 0 19 files changed, 143 insertions(+), 57 deletions(-) create mode 100644 packages/PyTorch/.gitignore create mode 100644 packages/PyTorch/LICENSE delete mode 100644 packages/PyTorch/pytorch_mr/__version__.py delete mode 100644 packages/PyTorch/pytorch_mr/tests/README.md delete mode 100644 packages/PyTorch/pytorch_mr/tests/__init__.py rename packages/PyTorch/{ => src}/pytorch_mr/__init__.py (100%) create mode 100644 packages/PyTorch/src/pytorch_mr/__version__.py rename packages/PyTorch/{ => src}/pytorch_mr/core.py (100%) create mode 100644 packages/PyTorch/tests/README.md rename packages/PyTorch/{pytorch_mr => }/tests/conftest.py (92%) rename packages/PyTorch/{pytorch_mr => }/tests/run_tests_by_chapter.py (84%) rename packages/PyTorch/{pytorch_mr => }/tests/test_ch1_3_numpy_parity.py (100%) rename packages/PyTorch/{pytorch_mr => }/tests/test_ch3.py (100%) rename packages/PyTorch/{pytorch_mr => }/tests/test_ch4_6.py (100%) rename packages/PyTorch/{pytorch_mr => }/tests/test_ch8.py (100%) rename packages/PyTorch/{pytorch_mr => }/tests/test_ch9_11.py (100%) 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 index 7bbbd21..c47f197 100644 --- a/packages/PyTorch/README.md +++ b/packages/PyTorch/README.md @@ -2,7 +2,12 @@ PyTorch implementation of the routines from *Modern Robotics: Mechanics, Planning, and Control*, aligned with the NumPy `modern_robotics` reference in `packages/Python`. -## Install +## Install (global) +```bash +pip install pytorch-mr +``` + +## Install (local / editable) From this directory (`packages/PyTorch`): @@ -10,7 +15,7 @@ From this directory (`packages/PyTorch`): pip install -e . ``` -Dependencies: `numpy`, `torch` (see `setup.py`). +Dependencies are declared in `pyproject.toml` (`numpy`, `torch`). ## Import @@ -22,29 +27,32 @@ from pytorch_mr.core import MatrixExp3, FKinBody ## Tests (NumPy parity) -Tests compare against the reference `modern_robotics` package. Install the NumPy library from the sibling directory or from PyPI: +Install the reference `modern_robotics` (local tree or PyPI) and `pytest`: ```bash -pip install -e ../Python # local reference, from packages/PyTorch +pip install -e ../Python # local reference # or pip install modern_robotics pip install pytest ``` -Run tests with both package roots on `PYTHONPATH` (from `packages/` in this repo): +From `packages/PyTorch` (uses `[tool.pytest.ini_options]` `pythonpath = ["src"]`): ```bash -cd packages -PYTHONPATH=PyTorch:Python python -m pytest PyTorch/pytorch_mr/tests -q +cd packages/PyTorch +python -m pytest tests -q ``` -Or use the chapter summary script: +From `packages/` with explicit `PYTHONPATH`: ```bash cd packages -PYTHONPATH=PyTorch:Python python PyTorch/pytorch_mr/tests/run_tests_by_chapter.py +PYTHONPATH=PyTorch/src:Python python -m pytest PyTorch/tests -q ``` -## Publishing to PyPI (optional) +Chapter summary: -Build and upload with standard `twine` / `build` workflows; ensure the package name `pytorch_mr` is available on your target index. +```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 index 4a85092..8ef6e38 100644 --- a/packages/PyTorch/pyproject.toml +++ b/packages/PyTorch/pyproject.toml @@ -1,3 +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/pytorch_mr/__version__.py b/packages/PyTorch/pytorch_mr/__version__.py deleted file mode 100644 index d930acc..0000000 --- a/packages/PyTorch/pytorch_mr/__version__.py +++ /dev/null @@ -1,2 +0,0 @@ - -__version__ = '1.1.0' diff --git a/packages/PyTorch/pytorch_mr/tests/README.md b/packages/PyTorch/pytorch_mr/tests/README.md deleted file mode 100644 index cd172a4..0000000 --- a/packages/PyTorch/pytorch_mr/tests/README.md +++ /dev/null @@ -1,25 +0,0 @@ -# Unit tests for `pytorch_mr` - -Tests live next to the PyTorch package under `packages/PyTorch/pytorch_mr/tests/`. - -They compare outputs to the NumPy reference `modern_robotics` (`packages/Python`). Ensure both are importable: - -- From repo `packages/` directory: - -```bash -PYTHONPATH=PyTorch:Python python -m pytest PyTorch/pytorch_mr/tests -q -``` - -- Chapter summary: - -```bash -PYTHONPATH=PyTorch:Python python PyTorch/pytorch_mr/tests/run_tests_by_chapter.py -``` - -After `pip install -e .` from `packages/PyTorch`, `pip install pytest`, and a NumPy reference (`pip install modern_robotics` or `pip install -e ../Python`), run: - -```bash -PYTHONPATH=/path/to/repo/packages/Python python -m pytest --pyargs pytorch_mr.tests -q -``` - -(Use `PYTHONPATH` to point at the local `modern_robotics` tree when you are not using the PyPI package.) diff --git a/packages/PyTorch/pytorch_mr/tests/__init__.py b/packages/PyTorch/pytorch_mr/tests/__init__.py deleted file mode 100644 index c919acf..0000000 --- a/packages/PyTorch/pytorch_mr/tests/__init__.py +++ /dev/null @@ -1 +0,0 @@ -# Test package for pytorch_mr (optional discovery via pytest --pyargs pytorch_mr.tests) diff --git a/packages/PyTorch/setup.py b/packages/PyTorch/setup.py index 7da9231..867027e 100644 --- a/packages/PyTorch/setup.py +++ b/packages/PyTorch/setup.py @@ -1,6 +1,6 @@ from setuptools import find_packages, setup -exec(open('pytorch_mr/__version__.py').read()) +exec(open('src/pytorch_mr/__version__.py').read()) long_description = """ # Modern Robotics (PyTorch) @@ -11,15 +11,14 @@ 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="Huan Weng, Mikhail Todes, Jarvis Schultz, Bill Hunt, Ayush Gaggar", - author_email="huanweng@u.northwestern.edu", + 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" ), @@ -44,4 +43,4 @@ "torch", ], platforms="Linux, Mac, Windows", -) +) \ No newline at end of file diff --git a/packages/PyTorch/pytorch_mr/__init__.py b/packages/PyTorch/src/pytorch_mr/__init__.py similarity index 100% rename from packages/PyTorch/pytorch_mr/__init__.py rename to packages/PyTorch/src/pytorch_mr/__init__.py 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/pytorch_mr/core.py b/packages/PyTorch/src/pytorch_mr/core.py similarity index 100% rename from packages/PyTorch/pytorch_mr/core.py rename to packages/PyTorch/src/pytorch_mr/core.py 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/pytorch_mr/tests/conftest.py b/packages/PyTorch/tests/conftest.py similarity index 92% rename from packages/PyTorch/pytorch_mr/tests/conftest.py rename to packages/PyTorch/tests/conftest.py index 03e4c39..a2848ec 100644 --- a/packages/PyTorch/pytorch_mr/tests/conftest.py +++ b/packages/PyTorch/tests/conftest.py @@ -8,11 +8,11 @@ _HERE = os.path.dirname(__file__) -# Prefer installed packages; fall back to sibling repo layout (packages/PyTorch + packages/Python). +# 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")) + _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 @@ -20,9 +20,9 @@ try: import pytorch_mr.core as mr_t # noqa: E402 except ImportError: - _PKG_TORCH_ROOT = os.path.abspath(os.path.join(_HERE, "..", "..")) - if _PKG_TORCH_ROOT not in sys.path: - sys.path.insert(0, _PKG_TORCH_ROOT) + _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 diff --git a/packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py b/packages/PyTorch/tests/run_tests_by_chapter.py similarity index 84% rename from packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py rename to packages/PyTorch/tests/run_tests_by_chapter.py index 6d8e58c..b4d2635 100644 --- a/packages/PyTorch/pytorch_mr/tests/run_tests_by_chapter.py +++ b/packages/PyTorch/tests/run_tests_by_chapter.py @@ -13,36 +13,38 @@ class ChapterRun: exit_code: int -def _collect_count(pytest_args, env): +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): +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__)) - # packages/ — parent of PyTorch/ and Python/ - packages_dir = os.path.abspath(os.path.join(here, "..", "..", "..")) - torch_pkg_root = os.path.join(packages_dir, "PyTorch") + # 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_pkg_root + sep + python_pkg_root + extra = torch_src + sep + python_pkg_root env["PYTHONPATH"] = extra + (sep + env["PYTHONPATH"] if env.get("PYTHONPATH") else "") chapters = [ @@ -59,8 +61,8 @@ def main(): for name, path in chapters: rel = os.path.relpath(path, packages_dir) - collected = _collect_count([rel], env) - exit_code = _run([rel], env) + 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) diff --git a/packages/PyTorch/pytorch_mr/tests/test_ch1_3_numpy_parity.py b/packages/PyTorch/tests/test_ch1_3_numpy_parity.py similarity index 100% rename from packages/PyTorch/pytorch_mr/tests/test_ch1_3_numpy_parity.py rename to packages/PyTorch/tests/test_ch1_3_numpy_parity.py diff --git a/packages/PyTorch/pytorch_mr/tests/test_ch3.py b/packages/PyTorch/tests/test_ch3.py similarity index 100% rename from packages/PyTorch/pytorch_mr/tests/test_ch3.py rename to packages/PyTorch/tests/test_ch3.py diff --git a/packages/PyTorch/pytorch_mr/tests/test_ch4_6.py b/packages/PyTorch/tests/test_ch4_6.py similarity index 100% rename from packages/PyTorch/pytorch_mr/tests/test_ch4_6.py rename to packages/PyTorch/tests/test_ch4_6.py diff --git a/packages/PyTorch/pytorch_mr/tests/test_ch8.py b/packages/PyTorch/tests/test_ch8.py similarity index 100% rename from packages/PyTorch/pytorch_mr/tests/test_ch8.py rename to packages/PyTorch/tests/test_ch8.py diff --git a/packages/PyTorch/pytorch_mr/tests/test_ch9_11.py b/packages/PyTorch/tests/test_ch9_11.py similarity index 100% rename from packages/PyTorch/pytorch_mr/tests/test_ch9_11.py rename to packages/PyTorch/tests/test_ch9_11.py