diff --git a/RBDReference.py b/RBDReference.py index d88358e..1f5a92d 100644 --- a/RBDReference.py +++ b/RBDReference.py @@ -21,7 +21,8 @@ class RBDReference( _CentroidalMixin, _RegressorMixin, ): - def __init__(self, robotObj, use_joint_dynamics=False): + def __init__(self, robotObj, use_joint_dynamics=False, friction_model="sign", + friction_width=0.03): """Initialize RBDReference with a robot object. Parameters @@ -37,6 +38,18 @@ def __init__(self, robotObj, use_joint_dynamics=False): authoritative-oracle-preserving default. The id/fd GRADIENTS follow this same flag: when True they add the damping diagonal (d(damping*qd)/dqd) to dc_dqd / qdd_dqd (friction's subgradient is 0). + friction_model : str, optional + ``"sign"`` (default) applies Coulomb friction as ``friction*sign(qd)``, + whose derivatives are zero almost everywhere. ``"tanh"`` applies the + smooth model ``friction*tanh(qd/width)``; its first derivative + ``(friction/width)*sech^2(qd/width)`` enters the inverse-dynamics + gradient and ``-(2*friction/width^2)*tanh(qd/width)*sech^2(qd/width)`` + enters the qd-qd block of the second-order derivatives. Only read when + ``use_joint_dynamics`` is True. + friction_width : float or array_like, optional + The width ``w`` of the tanh model in rad/s (or m/s): one scalar for + every velocity slot, or one value per velocity slot (length nv). + Default 0.03. Only read under ``friction_model="tanh"``. Returns ------- @@ -45,6 +58,20 @@ def __init__(self, robotObj, use_joint_dynamics=False): """ self.robot = robotObj # instance of Robot Object class created by URDFparser self.use_joint_dynamics = use_joint_dynamics + if friction_model not in ("sign", "tanh"): + raise ValueError( + f"friction_model must be 'sign' or 'tanh', got {friction_model!r}") + self.friction_model = friction_model + width = np.asarray(friction_width, dtype=np.float64) + if width.ndim == 0: + width = np.full(self.robot.get_num_vel(), float(width)) + if width.shape != (self.robot.get_num_vel(),): + raise ValueError( + "friction_width must be a scalar or one value per velocity slot " + f"(length {self.robot.get_num_vel()}); got shape {width.shape}") + if not np.all(width > 0.0): + raise ValueError("friction_width must be positive") + self.friction_width = width self._spatial_xmat_derivative_func_cache = {} self._spatial_xmat_second_derivative_func_cache = {} @@ -2599,11 +2626,13 @@ def inverse_dynamics_fpass(self, q, qd, qdd=None, GRAVITY=-9.81, f_ext=None): return (v, a, f) def _joint_dynamics_bias(self, qd): - """Per-v-slot joint-dynamics bias tau += damping*qd + friction*sign(qd). + """Per-v-slot joint-dynamics bias tau += damping*qd + friction*s(qd). Viscous damping and Coulomb friction are joint-local generalized forces that oppose motion. They are returned as an nv-vector added to the RNEA bias (inverse_dynamics) and subtracted from the available torque in ABA. + The friction shape s(qd) is sign(qd) under friction_model="sign" and + tanh(qd/w) under friction_model="tanh" (w = friction_width per slot). Returns the zero vector when no joint declares damping/friction (the common case), so this is numerically a no-op there. `qd` is the @@ -2631,9 +2660,65 @@ def _joint_dynamics_bias(self, qd): idx_list = idx if isinstance(idx, (list, tuple, np.ndarray)) else [idx] for k in idx_list: qd_k = qd[k] - bias[k] += alpha * (b * qd_k + fr * np.sign(qd_k)) + bias[k] += alpha * (b * qd_k + fr * self._friction_shape(qd_k, k)) return bias + def _friction_shape(self, qd_k, k): + """The friction velocity shape for slot k: sign(qd) or tanh(qd/w_k).""" + if self.friction_model == "tanh": + return np.tanh(qd_k / self.friction_width[k]) + return np.sign(qd_k) + + def _joint_dynamics_derivative_diag(self, qd, order): + """Per-v-slot derivative of the joint-dynamics bias with respect to the + slot's own velocity: the first (order=1) or second (order=2) derivative. + + order=1 is alpha*(damping + (friction/w) sech^2(qd/w)) under the tanh + model and alpha*damping under the sign model (the Coulomb term has zero + subgradient almost everywhere). order=2 is alpha*(-(2 friction/w^2) + tanh(qd/w) sech^2(qd/w)) under tanh and zero under sign (damping is + linear in qd). Mirrors _joint_dynamics_bias: iterate bodies, map each to + its reduced v-slot with the mimic multiplier, skip the floating root, and + accumulate so mimic joints sharing a slot sum. Zero when + use_joint_dynamics is False or no joint declares damping/friction. + """ + n = self.robot.get_num_vel() + diag = np.zeros(n) + if not self.use_joint_dynamics: + return diag + has_damping = self.robot.robot_has_joint_damping() + has_friction = self.robot.robot_has_joint_friction() + if not (has_damping or has_friction): + return diag + tanh_model = self.friction_model == "tanh" + if tanh_model and has_friction and qd is None: + raise ValueError("qd is required for the tanh friction derivatives") + for jid in range(self.robot.get_num_bodies()): + if self.robot.floating_base and self.robot.get_parent_id(jid) == -1: + continue # floating root carries no damping/friction + b = float(self.robot.get_damping_by_id(jid)) if has_damping else 0.0 + fr = float(self.robot.get_friction_by_id(jid)) if has_friction else 0.0 + if b == 0.0 and fr == 0.0: + continue + idx = self.robot.get_joint_index_v(jid) + alpha = self._mimic_multiplier(jid) + idx_list = idx if isinstance(idx, (list, tuple, np.ndarray)) else [idx] + for k in idx_list: + if order == 1: + val = b + else: + val = 0.0 + if tanh_model and fr != 0.0: + w = self.friction_width[k] + th = np.tanh(qd[k] / w) + sech2 = 1.0 - th * th + if order == 1: + val += (fr / w) * sech2 + else: + val += -(2.0 * fr / (w * w)) * th * sech2 + diag[k] += alpha * val + return diag + def inverse_dynamics_bpass(self, q, f): """Perform the backward pass of the Recursive Newton-Euler Algorithm. @@ -3825,7 +3910,7 @@ def inverse_dynamics_gradient_bpass_dq(self, q, f, df_dq): return dc_dq - def inverse_dynamics_gradient_bpass_dqd(self, q, df_dqd): + def inverse_dynamics_gradient_bpass_dqd(self, q, df_dqd, qd=None): """Backward pass gradient with respect to joint velocities for RNEA. Parameters @@ -3861,28 +3946,17 @@ def inverse_dynamics_gradient_bpass_dqd(self, q, df_dqd): df_dqd[:,:,parent_ind] += np.matmul(np.transpose(Xmat),df_dqd[:,:,ind]) - # Joint-local viscous damping contributes a diagonal term to dc_dqd: - # the value bias folds c += alpha*damping*qd_v into each joint's v-slot - # (see _joint_dynamics_bias), so d(c)/d(qd_v) = alpha*damping on the - # diagonal. Coulomb friction (alpha*f*sign(qd)) has zero (sub)gradient - # a.e., so it contributes NOTHING here -- gate on damping only. Mirrors - # _joint_dynamics_bias exactly: iterate bodies, map to the reduced - # v-slot via get_joint_index_v + the mimic multiplier, skip the floating - # root (it carries no damping), and ACCUMULATE so mimic joints sharing a - # slot sum. Gated on use_joint_dynamics (connected to the value path), - # NOT a separate kwarg. - if self.use_joint_dynamics and self.robot.robot_has_joint_damping(): - for jid in range(NB): - if self.robot.floating_base and self.robot.get_parent_id(jid) == -1: - continue # floating root carries no damping - b = float(self.robot.get_damping_by_id(jid)) - if b == 0.0: - continue - idx = self.robot.get_joint_index_v(jid) - alpha = self._mimic_multiplier(jid) - idx_list = idx if isinstance(idx, (list, tuple, np.ndarray)) else [idx] - for k in idx_list: - dc_dqd[k, k] += alpha * b + # The joint-dynamics bias contributes a diagonal term to dc_dqd: the + # value path folds c += alpha*(damping*qd_v + friction*s(qd_v)) into each + # joint's v-slot (see _joint_dynamics_bias), so d(c)/d(qd_v) on the + # diagonal is alpha*damping under the sign model (the Coulomb term has + # zero subgradient almost everywhere) and alpha*(damping + (friction/w) + # sech^2(qd_v/w)) under the tanh model. _joint_dynamics_derivative_diag + # mirrors the value fold (bodies to reduced v-slots with the mimic + # multiplier, floating root skipped, shared slots accumulate). Gated on + # use_joint_dynamics (connected to the value path), NOT a separate kwarg. + if self.use_joint_dynamics: + dc_dqd[np.diag_indices_from(dc_dqd)] += self._joint_dynamics_derivative_diag(qd, 1) return dc_dqd @@ -3947,7 +4021,7 @@ def inverse_dynamics_gradient( # backward pass, dqd (joint-damping gradient gated internally on # self.use_joint_dynamics, so it follows the value path automatically -- # forward_dynamics_gradient inherits it with no extra plumbing). - dc_dqd = self.inverse_dynamics_gradient_bpass_dqd(q, df_dqd) + dc_dqd = self.inverse_dynamics_gradient_bpass_dqd(q, df_dqd, qd=qd) if public_output: return self._denormalize_inverse_dynamics_gradient_output(dc_dq, dc_dqd) @@ -4600,6 +4674,7 @@ def subtree_vel_inds(subtree): if self.robot.floating_base: d2tau_dq = d2tau_dq + self._floating_gravity_d2tau_dq_lie_direct(q, GRAVITY) + self._add_joint_dynamics_so_diag(d2tau_dqd, qd) return d2tau_dq, d2tau_dqd, d2tau_dvdq, dM_dq def idsva_so_world_frame(self, q, qd, qdd, GRAVITY=-9.81): @@ -4847,8 +4922,24 @@ def idsva_so_world_frame(self, q, qd, qdd, GRAVITY=-9.81): d2tau_dvdq = np.einsum('ia,ijk,jb,kc->abc', R, d2tau_dvdq, R, R, optimize=True) dM_dq = np.einsum('ia,ijk,jb,kc->abc', R, dM_dq, R, R, optimize=True) + self._add_joint_dynamics_so_diag(d2tau_dqd, qd) return d2tau_dq, d2tau_dqd, d2tau_dvdq, dM_dq + def _add_joint_dynamics_so_diag(self, d2tau_dqd, qd): + """Add the joint-dynamics second derivative to the qd-qd block in place. + + The bias is linear in qd under the sign model (damping) with a zero + Coulomb subgradient, so this is a no-op there. Under the tanh model the + term -(2 friction/w^2) tanh(qd/w) sech^2(qd/w) lands on the triple + diagonal d2tau[k, k, k] of the reduced (public) velocity slots; it is + applied after any mimic fold, exactly where the CUDA inners add it. + """ + if not self.use_joint_dynamics: + return + so = self._joint_dynamics_derivative_diag(np.asarray(qd, dtype=np.float64), 2) + n = so.shape[0] + d2tau_dqd[np.arange(n), np.arange(n), np.arange(n)] += so + def idsva_so(self, q, qd, qdd, GRAVITY = -9.81): """Compute second-order derivatives of inverse dynamics. diff --git a/_plant.py b/_plant.py index 448668c..97fd6ac 100644 --- a/_plant.py +++ b/_plant.py @@ -403,6 +403,527 @@ def ee_pos_cost(self, q, p_des, W, ee=0): hess_x[:nv, :nv] = Jp.T @ (W[:, None] * Jp) # J_p^T diag(W) J_p return value, grad_x, hess_x + def gravity_torque(self, q, GRAVITY=-9.81): + """g(q) = RNEA(q, 0, 0), the generalized gravity (size nv); the oracle of + grid_plant::gravity_torque.""" + nv = self.robot.get_num_vel() + return np.asarray(self.inverse_dynamics(q, np.zeros(nv), np.zeros(nv), GRAVITY=GRAVITY)[0], dtype=np.float64).reshape(-1) + + # ----- cost menu: activations, activated terms, linear feature, named frames ----- + + ACTIVATION_QUADRATIC = 0 + ACTIVATION_PSEUDO_HUBER = 1 + + ACTIVATION_RELAXED_LOG_BARRIER = 2 + ACTIVATION_QUADRATIC_BARRIER = 3 + + @staticmethod + def _relaxed_log_barrier_1d(z, mu, delta): + """Grandia et al. relaxed log barrier on z >= 0: -mu log z above delta, the + quadratic extension below it (value and slope continuous at delta). + Returns (value, d/dz, d2/dz2) arrays. z may be +inf (no bound): zero.""" + z = np.asarray(z, dtype=np.float64) + val = np.zeros_like(z); d1 = np.zeros_like(z); d2 = np.zeros_like(z) + finite = np.isfinite(z) + above = finite & (z > delta) + below = finite & ~(z > delta) + val[above] = -mu * np.log(z[above]) + d1[above] = -mu / z[above] + d2[above] = mu / (z[above] ** 2) + t = (z[below] - 2.0 * delta) / delta + val[below] = mu * (0.5 * t * t - 0.5 - np.log(delta)) + d1[below] = mu * t / delta + d2[below] = mu / (delta * delta) * np.ones_like(t) + return val, d1, d2 + + @staticmethod + def _activation(kind, W, r, eps, v=None, lo=None, hi=None): + """Per-entry activation a(r; w, eps) with its first and second derivatives. + + kind 0 (quadratic): a = 1/2 w r^2, a' = w r, a'' = w. + kind 1 (pseudo-Huber): a = w (sqrt(r^2 + eps^2) - eps), a' = w r / s, + a'' = w eps^2 / s^3 with s = sqrt(r^2 + eps^2); eps must be positive. + kind 2 (relaxed log barrier): on the feature value v in [lo, hi], with + w = mu and eps = delta: B(v - lo) + B(hi - v) where B is -mu log z above + delta and its quadratic extension below; an infinite bound is zero. + kind 3 (quadratic barrier): 1/2 w (max(0, lo - v)^2 + max(0, v - hi)^2). + The barrier kinds read v, lo, hi; r is the residual the other kinds use. + Returns (value (M,), slope (M,), curvature (M,)), the derivatives with + respect to r (equal to those with respect to v: v = r + target). + """ + W = np.asarray(W, dtype=np.float64).reshape(-1) + r = np.asarray(r, dtype=np.float64).reshape(-1) + if kind == _PlantMixin.ACTIVATION_QUADRATIC: + return 0.5 * W * r * r, W * r, W.copy() + if kind == _PlantMixin.ACTIVATION_PSEUDO_HUBER: + eps = float(eps) + if not eps > 0.0: + raise ValueError("the pseudo-Huber activation needs eps > 0") + s = np.sqrt(r * r + eps * eps) + return W * (s - eps), W * r / s, W * eps * eps / (s * s * s) + if kind in (_PlantMixin.ACTIVATION_RELAXED_LOG_BARRIER, _PlantMixin.ACTIVATION_QUADRATIC_BARRIER): + if v is None or lo is None or hi is None: + raise ValueError("barrier activations need the feature value v and the bounds lo, hi") + v = np.asarray(v, dtype=np.float64).reshape(-1) + lo = np.asarray(lo, dtype=np.float64).reshape(-1) + hi = np.asarray(hi, dtype=np.float64).reshape(-1) + if kind == _PlantMixin.ACTIVATION_QUADRATIC_BARRIER: + below = np.where(np.isfinite(lo), np.maximum(0.0, lo - v), 0.0) + above = np.where(np.isfinite(hi), np.maximum(0.0, v - hi), 0.0) + value = 0.5 * W * (below * below + above * above) + slope = W * (above - below) + curv = W * ((below > 0) | (above > 0)).astype(np.float64) + return value, slope, curv + delta = float(eps) + if not delta > 0.0: + raise ValueError("the relaxed log barrier needs delta = eps > 0") + value = np.zeros_like(r); slope = np.zeros_like(r); curv = np.zeros_like(r) + for k in range(r.size): + mu = W[k] + if np.isfinite(lo[k]): + val, d1, d2 = _PlantMixin._relaxed_log_barrier_1d(np.array([v[k] - lo[k]]), mu, delta) + value[k] += val[0]; slope[k] += d1[0]; curv[k] += d2[0] + if np.isfinite(hi[k]): + val, d1, d2 = _PlantMixin._relaxed_log_barrier_1d(np.array([hi[k] - v[k]]), mu, delta) + value[k] += val[0]; slope[k] -= d1[0]; curv[k] += d2[0] + return value, slope, curv + raise ValueError(f"unknown activation kind {kind!r}") + + @staticmethod + def _skew3(v): + return np.array([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]]) + + @staticmethod + def _so3_log(R): + """phi = log(R) on the canonical branch |phi| <= pi (3-vector).""" + R = np.asarray(R, dtype=np.float64) + c = np.clip((np.trace(R) - 1.0) * 0.5, -1.0, 1.0) + theta = np.arccos(c) + w = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]]) + if theta < 1e-6: + return 0.5 * w + if theta < np.pi - 1e-4: + return theta / (2.0 * np.sin(theta)) * w + # near pi: axis from the symmetric part (R + I)/2 = a a^T, sign from w + S = 0.5 * (R + np.eye(3)) + k = int(np.argmax(np.diag(S))) + a = S[:, k] / np.sqrt(max(S[k, k], 1e-18)) + a = a / np.linalg.norm(a) + if np.dot(a, w) < 0.0: + a = -a + return theta * a + + @staticmethod + def _so3_right_jacobian_inv(phi): + """Jr(phi)^-1 = I + 1/2 [phi]x + (1/theta^2 - (1 + cos theta)/(2 theta sin theta)) [phi]x^2.""" + phi = np.asarray(phi, dtype=np.float64) + theta = np.linalg.norm(phi) + S = _PlantMixin._skew3(phi) + if theta < 1e-5: + coef = 1.0 / 12.0 + else: + coef = 1.0 / (theta * theta) - (1.0 + np.cos(theta)) / (2.0 * theta * np.sin(theta)) + return np.eye(3) + 0.5 * S + coef * (S @ S) + + def _frame_placement(self, q, frame_name): + """(R_f, p_f, target_id) of a named articulated or fixed joint frame, the + fixed joint's offset included (the frame_jacobian resolution).""" + Xw, _ = self._frame_world_placement_and_chain(q) + target_id, X_off = self._resolve_frame_joint(frame_name) + if target_id == -1: + X = X_off + else: + X = Xw[target_id] @ X_off + return X[:3, :3], X[:3, 3], target_id + + def frame_position_cost(self, q, frame_name, p_des, W, activation=0, eps=0.0, lo=None, hi=None): + """Position residual r = p_f(q) - p_des of a named frame under an activation. + Returns (value, grad_x (nx,) = [J_p^T a'(r); 0], GN hess_x (nx, nx)).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + R_f, p_f, _ = self._frame_placement(q, frame_name) + Jp = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL_WORLD_ALIGNED"))[:3, :] + value, slope, curv = self._activation(activation, W, p_f - np.asarray(p_des, dtype=np.float64).reshape(-1), eps, p_f, lo, hi) + grad_x = np.zeros(nx); grad_x[:nv] = Jp.T @ slope + hess_x = np.zeros((nx, nx)); hess_x[:nv, :nv] = Jp.T @ (curv[:, None] * Jp) + return float(np.sum(value)), grad_x, hess_x + + def frame_orientation_cost(self, q, frame_name, R_des, W, activation=0, eps=0.0): + """Log-map orientation residual e = log(R_des^T R_f(q)) (3-vector) under an + activation. With the frame's body-axes angular Jacobian J_w (LOCAL rows 3..5) + the residual Jacobian is J_e = Jr(e)^-1 J_w. Returns (value, grad_x, GN hess_x).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + R_f, _, _ = self._frame_placement(q, frame_name) + R_des = np.asarray(R_des, dtype=np.float64).reshape(3, 3, order="F") + e = self._so3_log(R_des.T @ R_f) + Jw = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL"))[3:, :] + Je = self._so3_right_jacobian_inv(e) @ Jw + value, slope, curv = self._activation(activation, W, e, eps) + grad_x = np.zeros(nx); grad_x[:nv] = Je.T @ slope + hess_x = np.zeros((nx, nx)); hess_x[:nv, :nv] = Je.T @ (curv[:, None] * Je) + return float(np.sum(value)), grad_x, hess_x + + def _frame_velocity_q_derivative(self, q, qd, frame_name): + """d(J(q) qd)/dq for the LOCAL_WORLD_ALIGNED frame Jacobian (6 x nv), from + the same world-axis column construction frame_jacobian uses. Fixed base.""" + if self.robot.floating_base: + raise NotImplementedError("frame velocity residual: fixed base only in this release") + nv = self.robot.get_num_vel() + Xw, _ = self._frame_world_placement_and_chain(q) + target_id, X_off = self._resolve_frame_joint(frame_name) + p_f = (Xw[target_id] @ X_off)[:3, 3] + chain = sorted(self.robot.get_ancestors_by_id(target_id)) + [target_id] + + def column_info(j): + S = np.asarray(self.robot.get_S_by_id(j), dtype=np.float64) + if S.ndim == 1: + S = S.reshape(-1, 1) + R_j = Xw[j][:3, :3]; p_j = Xw[j][:3, 3] + inds = self.robot.get_joint_index_v(j) + inds = list(inds) if isinstance(inds, (list, tuple, np.ndarray)) else [inds] + joint = self.robot.get_joint_by_id(j) + scale = joint.get_mimic_multiplier() if getattr(joint, "is_mimic", False) else 1.0 + out = [] + for c in range(S.shape[1]): + vi = inds[c] if c < len(inds) else inds[-1] + rot = np.linalg.norm(S[:3, c]) > 0.5 + axis = R_j @ (S[:3, c] if rot else S[3:6, c]) + out.append((vi, rot, axis, p_j, scale)) + return out + + cols = {j: column_info(j) for j in chain} + dJqd = np.zeros((6, nv)) + qd = np.asarray(qd, dtype=np.float64).reshape(-1) + # Column j is a_j x (p_f - p_j) (revolute) or l_j (prismatic) over its axis + # a_j / l_j and origin p_j. A joint m moves, rigidly, everything downstream + # of it: for m an ancestor-or-self of j that is the axis, the origin and the + # frame point; for m downstream of j only the frame point p_f moves. + for j in chain: + for (vi, rot, a_j, p_j, s_j) in cols[j]: + w = s_j * qd[vi] + for m in chain: + for (vm, rot_m, a_m, p_m, s_m) in cols[m]: + if m <= j: + if rot and rot_m: + d_axis = np.cross(a_m, a_j) + d_lin = np.cross(d_axis, p_f - p_j) + np.cross(a_j, np.cross(a_m, p_f - p_j)) + dJqd[:3, vm] += s_m * w * d_lin + dJqd[3:, vm] += s_m * w * d_axis + elif (not rot) and rot_m: + dJqd[:3, vm] += s_m * w * np.cross(a_m, a_j) + # a prismatic m shifts the frame and the joint origins together: no change + elif rot: + # m downstream of j: only p_f moves, by a_m x (p_f - p_m) or l_m + d_pf = np.cross(a_m, p_f - p_m) if rot_m else a_m + dJqd[:3, vm] += s_m * w * np.cross(a_j, d_pf) + return dJqd + + def frame_velocity_cost(self, q, qd, frame_name, v_des, W, activation=0, eps=0.0, lo=None, hi=None): + """Frame velocity residual r = J(q) qd - v_des over the LOCAL_WORLD_ALIGNED + twist [linear; angular] (6) under an activation. The residual Jacobian over + x = [q; qd] is [d(J qd)/dq | J]. Returns (value, grad_x (nx,), GN hess_x).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + J = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL_WORLD_ALIGNED")) + qd = np.asarray(qd, dtype=np.float64).reshape(-1) + r = J @ qd - np.asarray(v_des, dtype=np.float64).reshape(-1) + Jx = np.hstack([self._frame_velocity_q_derivative(q, qd, frame_name), J]) + value, slope, curv = self._activation(activation, W, r, eps, J @ qd, lo, hi) + return float(np.sum(value)), Jx.T @ slope, Jx.T @ (curv[:, None] * Jx) + + @classmethod + def _activated_cost(cls, var, des, W, activation, eps, lo=None, hi=None): + var = np.asarray(var, dtype=np.float64).reshape(-1) + des = np.asarray(des, dtype=np.float64).reshape(-1) + value, slope, curv = cls._activation(activation, W, var - des, eps, var, lo, hi) + return float(np.sum(value)), slope, np.diag(curv) + + def state_cost(self, x, x_des, Q, activation=0, eps=0.0, lo=None, hi=None): + """sum_i a(x_i - x_des_i; Q_i, eps) over x = [q; qd]. The quadratic + activation is quadratic_state_cost; lo / hi are the barrier bounds on x. + Returns (value, grad (nx,), GN hess (nx, nx)).""" + return self._activated_cost(x, x_des, Q, activation, eps, lo, hi) + + def input_cost(self, u, u_des, R, activation=0, eps=0.0, lo=None, hi=None): + """sum_i a(u_i - u_des_i; R_i, eps). Returns (value, grad (nu,), GN hess (nu, nu)).""" + return self._activated_cost(u, u_des, R, activation, eps, lo, hi) + + def linear_feature_cost(self, v, A, b, target, W, activation=0, eps=0.0, lo=None, hi=None): + """sum_m a(r_m; W_m, eps) with r = A v + b - target, A of shape (M, N); lo / hi + are barrier bounds on the feature value A v + b. + + Returns (value, grad (N,) = A^T a'(r), GN hess (N, N) = A^T diag(a''(r)) A). + """ + A = np.asarray(A, dtype=np.float64) + v = np.asarray(v, dtype=np.float64).reshape(-1) + f = A @ v + np.asarray(b, dtype=np.float64).reshape(-1) + r = f - np.asarray(target, dtype=np.float64).reshape(-1) + value, slope, curv = self._activation(activation, W, r, eps, f, lo, hi) + return float(np.sum(value)), A.T @ slope, A.T @ (curv[:, None] * A) + + ACTIVATION_RELAXED_LOG_BARRIER = 2 + ACTIVATION_QUADRATIC_BARRIER = 3 + + @staticmethod + def _relaxed_log_barrier_1d(z, mu, delta): + """Grandia et al. relaxed log barrier on z >= 0: -mu log z above delta, the + quadratic extension below it (value and slope continuous at delta). + Returns (value, d/dz, d2/dz2) arrays. z may be +inf (no bound): zero.""" + z = np.asarray(z, dtype=np.float64) + val = np.zeros_like(z); d1 = np.zeros_like(z); d2 = np.zeros_like(z) + finite = np.isfinite(z) + above = finite & (z > delta) + below = finite & ~(z > delta) + val[above] = -mu * np.log(z[above]) + d1[above] = -mu / z[above] + d2[above] = mu / (z[above] ** 2) + t = (z[below] - 2.0 * delta) / delta + val[below] = mu * (0.5 * t * t - 0.5 - np.log(delta)) + d1[below] = mu * t / delta + d2[below] = mu / (delta * delta) * np.ones_like(t) + return val, d1, d2 + + @staticmethod + def _activation(kind, W, r, eps, v=None, lo=None, hi=None): + """Per-entry activation a(r; w, eps) with its first and second derivatives. + + kind 0 (quadratic): a = 1/2 w r^2, a' = w r, a'' = w. + kind 1 (pseudo-Huber): a = w (sqrt(r^2 + eps^2) - eps), a' = w r / s, + a'' = w eps^2 / s^3 with s = sqrt(r^2 + eps^2); eps must be positive. + kind 2 (relaxed log barrier): on the feature value v in [lo, hi], with + w = mu and eps = delta: B(v - lo) + B(hi - v) where B is -mu log z above + delta and its quadratic extension below; an infinite bound is zero. + kind 3 (quadratic barrier): 1/2 w (max(0, lo - v)^2 + max(0, v - hi)^2). + The barrier kinds read v, lo, hi; r is the residual the other kinds use. + Returns (value (M,), slope (M,), curvature (M,)), the derivatives with + respect to r (equal to those with respect to v: v = r + target). + """ + W = np.asarray(W, dtype=np.float64).reshape(-1) + r = np.asarray(r, dtype=np.float64).reshape(-1) + if kind == _PlantMixin.ACTIVATION_QUADRATIC: + return 0.5 * W * r * r, W * r, W.copy() + if kind == _PlantMixin.ACTIVATION_PSEUDO_HUBER: + eps = float(eps) + if not eps > 0.0: + raise ValueError("the pseudo-Huber activation needs eps > 0") + s = np.sqrt(r * r + eps * eps) + return W * (s - eps), W * r / s, W * eps * eps / (s * s * s) + if kind in (_PlantMixin.ACTIVATION_RELAXED_LOG_BARRIER, _PlantMixin.ACTIVATION_QUADRATIC_BARRIER): + if v is None or lo is None or hi is None: + raise ValueError("barrier activations need the feature value v and the bounds lo, hi") + v = np.asarray(v, dtype=np.float64).reshape(-1) + lo = np.asarray(lo, dtype=np.float64).reshape(-1) + hi = np.asarray(hi, dtype=np.float64).reshape(-1) + if kind == _PlantMixin.ACTIVATION_QUADRATIC_BARRIER: + below = np.where(np.isfinite(lo), np.maximum(0.0, lo - v), 0.0) + above = np.where(np.isfinite(hi), np.maximum(0.0, v - hi), 0.0) + value = 0.5 * W * (below * below + above * above) + slope = W * (above - below) + curv = W * ((below > 0) | (above > 0)).astype(np.float64) + return value, slope, curv + delta = float(eps) + if not delta > 0.0: + raise ValueError("the relaxed log barrier needs delta = eps > 0") + value = np.zeros_like(r); slope = np.zeros_like(r); curv = np.zeros_like(r) + for k in range(r.size): + mu = W[k] + if np.isfinite(lo[k]): + val, d1, d2 = _PlantMixin._relaxed_log_barrier_1d(np.array([v[k] - lo[k]]), mu, delta) + value[k] += val[0]; slope[k] += d1[0]; curv[k] += d2[0] + if np.isfinite(hi[k]): + val, d1, d2 = _PlantMixin._relaxed_log_barrier_1d(np.array([hi[k] - v[k]]), mu, delta) + value[k] += val[0]; slope[k] -= d1[0]; curv[k] += d2[0] + return value, slope, curv + raise ValueError(f"unknown activation kind {kind!r}") + + @staticmethod + def _skew3(v): + return np.array([[0.0, -v[2], v[1]], [v[2], 0.0, -v[0]], [-v[1], v[0], 0.0]]) + + @staticmethod + def _so3_log(R): + """phi = log(R) on the canonical branch |phi| <= pi (3-vector).""" + R = np.asarray(R, dtype=np.float64) + c = np.clip((np.trace(R) - 1.0) * 0.5, -1.0, 1.0) + theta = np.arccos(c) + w = np.array([R[2, 1] - R[1, 2], R[0, 2] - R[2, 0], R[1, 0] - R[0, 1]]) + if theta < 1e-6: + return 0.5 * w + if theta < np.pi - 1e-4: + return theta / (2.0 * np.sin(theta)) * w + # near pi: axis from the symmetric part (R + I)/2 = a a^T, sign from w + S = 0.5 * (R + np.eye(3)) + k = int(np.argmax(np.diag(S))) + a = S[:, k] / np.sqrt(max(S[k, k], 1e-18)) + a = a / np.linalg.norm(a) + if np.dot(a, w) < 0.0: + a = -a + return theta * a + + @staticmethod + def _so3_right_jacobian_inv(phi): + """Jr(phi)^-1 = I + 1/2 [phi]x + (1/theta^2 - (1 + cos theta)/(2 theta sin theta)) [phi]x^2.""" + phi = np.asarray(phi, dtype=np.float64) + theta = np.linalg.norm(phi) + S = _PlantMixin._skew3(phi) + if theta < 1e-5: + coef = 1.0 / 12.0 + else: + coef = 1.0 / (theta * theta) - (1.0 + np.cos(theta)) / (2.0 * theta * np.sin(theta)) + return np.eye(3) + 0.5 * S + coef * (S @ S) + + def _frame_placement(self, q, frame_name): + """(R_f, p_f, target_id) of a named articulated or fixed joint frame, the + fixed joint's offset included (the frame_jacobian resolution).""" + Xw, _ = self._frame_world_placement_and_chain(q) + target_id, X_off = self._resolve_frame_joint(frame_name) + if target_id == -1: + X = X_off + else: + X = Xw[target_id] @ X_off + return X[:3, :3], X[:3, 3], target_id + + def frame_position_cost(self, q, frame_name, p_des, W, activation=0, eps=0.0, lo=None, hi=None): + """Position residual r = p_f(q) - p_des of a named frame under an activation. + Returns (value, grad_x (nx,) = [J_p^T a'(r); 0], GN hess_x (nx, nx)).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + R_f, p_f, _ = self._frame_placement(q, frame_name) + Jp = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL_WORLD_ALIGNED"))[:3, :] + value, slope, curv = self._activation(activation, W, p_f - np.asarray(p_des, dtype=np.float64).reshape(-1), eps, p_f, lo, hi) + grad_x = np.zeros(nx); grad_x[:nv] = Jp.T @ slope + hess_x = np.zeros((nx, nx)); hess_x[:nv, :nv] = Jp.T @ (curv[:, None] * Jp) + return float(np.sum(value)), grad_x, hess_x + + def frame_orientation_cost(self, q, frame_name, R_des, W, activation=0, eps=0.0): + """Log-map orientation residual e = log(R_des^T R_f(q)) (3-vector) under an + activation. With the frame's body-axes angular Jacobian J_w (LOCAL rows 3..5) + the residual Jacobian is J_e = Jr(e)^-1 J_w. Returns (value, grad_x, GN hess_x).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + R_f, _, _ = self._frame_placement(q, frame_name) + R_des = np.asarray(R_des, dtype=np.float64).reshape(3, 3, order="F") + e = self._so3_log(R_des.T @ R_f) + Jw = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL"))[3:, :] + Je = self._so3_right_jacobian_inv(e) @ Jw + value, slope, curv = self._activation(activation, W, e, eps) + grad_x = np.zeros(nx); grad_x[:nv] = Je.T @ slope + hess_x = np.zeros((nx, nx)); hess_x[:nv, :nv] = Je.T @ (curv[:, None] * Je) + return float(np.sum(value)), grad_x, hess_x + + def _frame_velocity_q_derivative(self, q, qd, frame_name): + """d(J(q) qd)/dq for the LOCAL_WORLD_ALIGNED frame Jacobian (6 x nv), from + the same world-axis column construction frame_jacobian uses. Fixed base.""" + if self.robot.floating_base: + raise NotImplementedError("frame velocity residual: fixed base only in this release") + nv = self.robot.get_num_vel() + Xw, _ = self._frame_world_placement_and_chain(q) + target_id, X_off = self._resolve_frame_joint(frame_name) + p_f = (Xw[target_id] @ X_off)[:3, 3] + chain = sorted(self.robot.get_ancestors_by_id(target_id)) + [target_id] + + def column_info(j): + S = np.asarray(self.robot.get_S_by_id(j), dtype=np.float64) + if S.ndim == 1: + S = S.reshape(-1, 1) + R_j = Xw[j][:3, :3]; p_j = Xw[j][:3, 3] + inds = self.robot.get_joint_index_v(j) + inds = list(inds) if isinstance(inds, (list, tuple, np.ndarray)) else [inds] + joint = self.robot.get_joint_by_id(j) + scale = joint.get_mimic_multiplier() if getattr(joint, "is_mimic", False) else 1.0 + out = [] + for c in range(S.shape[1]): + vi = inds[c] if c < len(inds) else inds[-1] + rot = np.linalg.norm(S[:3, c]) > 0.5 + axis = R_j @ (S[:3, c] if rot else S[3:6, c]) + out.append((vi, rot, axis, p_j, scale)) + return out + + cols = {j: column_info(j) for j in chain} + dJqd = np.zeros((6, nv)) + qd = np.asarray(qd, dtype=np.float64).reshape(-1) + # Column j is a_j x (p_f - p_j) (revolute) or l_j (prismatic) over its axis + # a_j / l_j and origin p_j. A joint m moves, rigidly, everything downstream + # of it: for m an ancestor-or-self of j that is the axis, the origin and the + # frame point; for m downstream of j only the frame point p_f moves. + for j in chain: + for (vi, rot, a_j, p_j, s_j) in cols[j]: + w = s_j * qd[vi] + for m in chain: + for (vm, rot_m, a_m, p_m, s_m) in cols[m]: + if m <= j: + if rot and rot_m: + d_axis = np.cross(a_m, a_j) + d_lin = np.cross(d_axis, p_f - p_j) + np.cross(a_j, np.cross(a_m, p_f - p_j)) + dJqd[:3, vm] += s_m * w * d_lin + dJqd[3:, vm] += s_m * w * d_axis + elif (not rot) and rot_m: + dJqd[:3, vm] += s_m * w * np.cross(a_m, a_j) + # a prismatic m shifts the frame and the joint origins together: no change + elif rot: + # m downstream of j: only p_f moves, by a_m x (p_f - p_m) or l_m + d_pf = np.cross(a_m, p_f - p_m) if rot_m else a_m + dJqd[:3, vm] += s_m * w * np.cross(a_j, d_pf) + return dJqd + + def frame_velocity_cost(self, q, qd, frame_name, v_des, W, activation=0, eps=0.0, lo=None, hi=None): + """Frame velocity residual r = J(q) qd - v_des over the LOCAL_WORLD_ALIGNED + twist [linear; angular] (6) under an activation. The residual Jacobian over + x = [q; qd] is [d(J qd)/dq | J]. Returns (value, grad_x (nx,), GN hess_x).""" + nq = self.robot.get_num_pos(); nv = self.robot.get_num_vel(); nx = nq + nv + J = np.asarray(self.frame_jacobian(q, frame_name, "LOCAL_WORLD_ALIGNED")) + qd = np.asarray(qd, dtype=np.float64).reshape(-1) + r = J @ qd - np.asarray(v_des, dtype=np.float64).reshape(-1) + Jx = np.hstack([self._frame_velocity_q_derivative(q, qd, frame_name), J]) + value, slope, curv = self._activation(activation, W, r, eps, J @ qd, lo, hi) + return float(np.sum(value)), Jx.T @ slope, Jx.T @ (curv[:, None] * Jx) + + @classmethod + def _activated_cost(cls, var, des, W, activation, eps, lo=None, hi=None): + var = np.asarray(var, dtype=np.float64).reshape(-1) + des = np.asarray(des, dtype=np.float64).reshape(-1) + value, slope, curv = cls._activation(activation, W, var - des, eps, var, lo, hi) + return float(np.sum(value)), slope, np.diag(curv) + + def state_cost(self, x, x_des, Q, activation=0, eps=0.0, lo=None, hi=None): + """sum_i a(x_i - x_des_i; Q_i, eps) over x = [q; qd]. The quadratic + activation is quadratic_state_cost; lo / hi are the barrier bounds on x. + Returns (value, grad (nx,), GN hess (nx, nx)).""" + return self._activated_cost(x, x_des, Q, activation, eps, lo, hi) + + def input_cost(self, u, u_des, R, activation=0, eps=0.0, lo=None, hi=None): + """sum_i a(u_i - u_des_i; R_i, eps). Returns (value, grad (nu,), GN hess (nu, nu)).""" + return self._activated_cost(u, u_des, R, activation, eps, lo, hi) + + def linear_feature_cost(self, v, A, b, target, W, activation=0, eps=0.0, lo=None, hi=None): + """sum_m a(r_m; W_m, eps) with r = A v + b - target, A of shape (M, N); lo / hi + are barrier bounds on the feature value A v + b. + + Returns (value, grad (N,) = A^T a'(r), GN hess (N, N) = A^T diag(a''(r)) A). + """ + A = np.asarray(A, dtype=np.float64) + v = np.asarray(v, dtype=np.float64).reshape(-1) + f = A @ v + np.asarray(b, dtype=np.float64).reshape(-1) + r = f - np.asarray(target, dtype=np.float64).reshape(-1) + value, slope, curv = self._activation(activation, W, r, eps, f, lo, hi) + return float(np.sum(value)), A.T @ slope, A.T @ (curv[:, None] * A) + + def ee_pos_cost_activated(self, q, p_des, W, activation=0, eps=0.0, ee=0): + """End-effector position residual r = p(q) - p_des under an activation. + + Returns (value, grad_x (nx,) = [J_p^T a'(r); 0], GN hess_x (nx, nx) with the + q-block J_p^T diag(a''(r)) J_p). The quadratic activation is ee_pos_cost. + """ + nq = self.robot.get_num_pos() + nv = self.robot.get_num_vel() + nx = nq + nv + target = self._ee_target_name(ee) + p_des = np.asarray(p_des, dtype=np.float64).reshape(-1) + pose = np.asarray(self.end_effector_pose(q, ee_joint_names=target)[0], dtype=np.float64).reshape(-1) + Jp = np.asarray(self.end_effector_pose_gradient(q, ee_joint_names=target)[0], dtype=np.float64)[:3, :] + value, slope, curv = self._activation(activation, W, pose[:3] - p_des, eps) + grad_x = np.zeros(nx, dtype=np.float64) + grad_x[:nv] = Jp.T @ slope + hess_x = np.zeros((nx, nx), dtype=np.float64) + hess_x[:nv, :nv] = Jp.T @ (curv[:, None] * Jp) + return float(np.sum(value)), grad_x, hess_x + # ----- CoM-tracking cost (value, grad_x, GN hess_x) ----- def com_cost(self, q, p_des, W): diff --git a/tests/test_cost_menu_reference.py b/tests/test_cost_menu_reference.py new file mode 100644 index 0000000..6a23001 --- /dev/null +++ b/tests/test_cost_menu_reference.py @@ -0,0 +1,221 @@ +"""Cost menu reference checks: activations, activated terms, linear feature, named frames. + +GRiD-defined quantities (the CUDA grid_plant surface), so there is no Pinocchio +oracle; each is checked against a closed form, a central-difference gradient and a +Gauss-Newton Hessian recompute on iiwa14 (fixed) and go2 (floating): + - the quadratic activation reproduces quadratic_state_cost / quadratic_input_cost / + ee_pos_cost exactly, + - the pseudo-Huber activation has the stated value, slope and curvature and the + weighted-absolute-value and quadratic limits, + - state_cost / input_cost / linear_feature_cost / frame_position_cost gradients + match central differences and the Hessians are J^T diag(a'') J. +""" +from pathlib import Path + +import numpy as np +import pytest + +from URDFParser import URDFParser +from RBDReference import RBDReference + +_ASSETS = Path(__file__).resolve().parents[1] / "robot_assets" +_H = 1e-6 + + +def _ref(name, floating): + return RBDReference(URDFParser().parse(str(_ASSETS / f"{name}.urdf"), floating_base=floating)) + + +def _central(f, x, h=_H): + x = np.asarray(x, dtype=np.float64) + g = np.zeros_like(x) + for k in range(x.size): + xp = x.copy(); xp[k] += h + xm = x.copy(); xm[k] -= h + g[k] = (f(xp) - f(xm)) / (2.0 * h) + return g + + +def test_activation_closed_forms_and_limits(): + W = np.array([2.0, 3.0, 0.5]) + r = np.array([-0.4, 0.0, 0.7]) + v, g, c = RBDReference._activation(0, W, r, 0.0) + assert np.allclose(v, 0.5 * W * r * r) and np.allclose(g, W * r) and np.allclose(c, W) + eps = 0.05 + s = np.sqrt(r * r + eps * eps) + v, g, c = RBDReference._activation(1, W, r, eps) + assert np.allclose(v, W * (s - eps)) and np.allclose(g, W * r / s) and np.allclose(c, W * eps ** 2 / s ** 3) + # the slope is the derivative of the value and the curvature the derivative of the slope + for k in range(3): + fv = lambda rr: RBDReference._activation(1, W[k:k + 1], rr, eps)[0][0] + fg = lambda rr: RBDReference._activation(1, W[k:k + 1], rr, eps)[1][0] + assert abs(_central(fv, r[k:k + 1])[0] - g[k]) < 1e-6 + assert abs(_central(fg, r[k:k + 1])[0] - c[k]) < 1e-5 + # limits: tiny eps is the weighted absolute value; |r| << eps is the quadratic w r^2 / (2 eps) + v_small, _, _ = RBDReference._activation(1, W, r, 1e-9) + assert np.allclose(v_small, W * np.abs(r), atol=1e-8) + v_big, _, c_big = RBDReference._activation(1, W, 1e-4 * r, 1.0) + assert np.allclose(v_big, 0.5 * W * (1e-4 * r) ** 2, rtol=1e-6) and np.allclose(c_big, W, rtol=1e-6) + with pytest.raises(ValueError): + RBDReference._activation(1, W, r, 0.0) + with pytest.raises(ValueError): + RBDReference._activation(7, W, r, 0.1) + + +@pytest.mark.parametrize("name,floating", [("iiwa14", False), ("go2", True)], ids=["iiwa14", "go2"]) +def test_activated_terms_against_finite_differences(name, floating): + ref = _ref(name, floating) + nq, nv = ref.robot.get_num_pos(), ref.robot.get_num_vel() + nx = nq + nv + rng = np.random.default_rng(5) + q = ref.robot.normalize_floating_base_q_input(rng.standard_normal(nq)) if floating else rng.standard_normal(nq) + qd = rng.standard_normal(nv) + u = rng.standard_normal(nv) + x = np.concatenate([q, qd]) + x_des = 0.1 * np.arange(nx); Q = 1.0 + 0.5 * np.arange(nx) + u_des = -0.05 * np.arange(nv); R = 2.0 + 0.1 * np.arange(nv) + eps = 0.03 + # quadratic activation reproduces the quadratic terms exactly + for a, b in zip(ref.state_cost(x, x_des, Q, 0, 0.0), ref.quadratic_state_cost(x, x_des, Q)): + assert np.array_equal(np.asarray(a), np.asarray(b)) + for a, b in zip(ref.input_cost(u, u_des, R, 0, 0.0), ref.quadratic_input_cost(u, u_des, R)): + assert np.array_equal(np.asarray(a), np.asarray(b)) + # pseudo-Huber state and input: gradient vs FD of the value, Hessian diagonal of a'' + val, grad, hess = ref.state_cost(x, x_des, Q, 1, eps) + assert np.allclose(grad, _central(lambda xx: ref.state_cost(xx, x_des, Q, 1, eps)[0], x), atol=1e-6) + assert np.allclose(np.diag(hess), ref._activation(1, Q, x - x_des, eps)[2]) and np.count_nonzero(hess - np.diag(np.diag(hess))) == 0 + val, grad, hess = ref.input_cost(u, u_des, R, 1, eps) + assert np.allclose(grad, _central(lambda uu: ref.input_cost(uu, u_des, R, 1, eps)[0], u), atol=1e-6) + # linear feature on x (2 rows) and on u (1 row) + A = rng.standard_normal((2, nx)); b = np.array([0.0, 0.3]); t = np.array([0.2, 0.3]); W = np.array([4.0, 5.0]) + val, grad, hess = ref.linear_feature_cost(x, A, b, t, W, 1, eps) + assert np.allclose(grad, _central(lambda xx: ref.linear_feature_cost(xx, A, b, t, W, 1, eps)[0], x), atol=1e-6) + curv = ref._activation(1, W, A @ x + b - t, eps)[2] + assert np.allclose(hess, A.T @ (curv[:, None] * A)) + Au = rng.standard_normal((1, nv)) + val, grad, hess = ref.linear_feature_cost(u, Au, [0.1], [-0.1], [5.0], 1, eps) + assert np.allclose(grad, _central(lambda uu: ref.linear_feature_cost(uu, Au, [0.1], [-0.1], [5.0], 1, eps)[0], u), atol=1e-6) + # the quadratic activation on a linear feature is 1/2 (A x + b - t)^T diag(W) (A x + b - t) + rr = A @ x + b - t + assert np.isclose(ref.linear_feature_cost(x, A, b, t, W, 0, 0.0)[0], 0.5 * np.sum(W * rr * rr)) + + +@pytest.mark.parametrize("name,floating", [("iiwa14", False), ("go2", True)], ids=["iiwa14", "go2"]) +def test_frame_zero_is_the_end_effector_cost(name, floating): + """frame_position_cost on the end-effector joint under the quadratic activation is + ee_pos_cost; under pseudo-Huber it carries the Gauss-Newton structure from the + position Jacobian and zero qd blocks (the floating go2 included).""" + ref = _ref(name, floating) + nq, nv = ref.robot.get_num_pos(), ref.robot.get_num_vel() + rng = np.random.default_rng(9) + q = ref.robot.normalize_floating_base_q_input(rng.standard_normal(nq)) if floating else rng.standard_normal(nq) + frame = ref.robot.get_joint_by_id(ref.robot.get_leaf_nodes()[0]).get_name() + W = np.array([10.0, 11.0, 12.0]); p_des = np.array([0.1, -0.2, 0.3]); eps = 0.02 + for a, b in zip(ref.frame_position_cost(q, frame, p_des, W, 0, 0.0), ref.ee_pos_cost(q, p_des, W)): + assert np.allclose(np.asarray(a), np.asarray(b), atol=1e-12) + val, grad_x, hess_x = ref.frame_position_cost(q, frame, p_des, W, 1, eps) + assert np.count_nonzero(grad_x[nv:]) == 0 and np.count_nonzero(hess_x[nv:, :]) == 0 + Jp = np.asarray(ref.frame_jacobian(q, frame))[:3, :] + _, p, _ = ref._frame_placement(q, frame) + _, slope, curv = ref._activation(1, W, p - p_des, eps) + assert np.allclose(grad_x[:nv], Jp.T @ slope) and np.allclose(hess_x[:nv, :nv], Jp.T @ (curv[:, None] * Jp)) + + +# ---- named-frame residuals and barrier activations ---- + +def _frames(ref): + robot = ref.robot + return [robot.get_joint_by_id(robot.get_leaf_nodes()[0]).get_name(), robot.get_joint_by_id(3).get_name()] + + +@pytest.mark.parametrize("name,floating", [("iiwa14", False), ("fr3", False)], ids=["iiwa14", "fr3"]) +def test_frame_position_and_orientation_costs(name, floating): + ref = _ref(name, floating) + nq, nv = ref.robot.get_num_pos(), ref.robot.get_num_vel() + rng = np.random.default_rng(11) + q = rng.standard_normal(nq) + W = np.array([2.0, 3.0, 4.0]); eps = 0.05 + for frame in _frames(ref): + p_des = np.array([0.1, 0.2, 0.3]) + val, grad, hess = ref.frame_position_cost(q, frame, p_des, W, 1, eps) + fd = _central(lambda qq: ref.frame_position_cost(qq, frame, p_des, W, 1, eps)[0], q) + assert np.allclose(grad[:nv], fd, atol=1e-6) and np.count_nonzero(grad[nv:]) == 0 + # orientation: a target rotated from the current frame by phi recovers phi as the residual + R_f, _, _ = ref._frame_placement(q, frame) + phi = np.array([0.3, -0.2, 0.5]) + R_des = R_f @ ref._so3_right_jacobian_inv(np.zeros(3)) # identity start + from scipy.spatial.transform import Rotation + R_des = R_f @ Rotation.from_rotvec(-phi).as_matrix() + assert np.allclose(ref._so3_log(R_des.T @ R_f), phi, atol=1e-12) + val, grad, hess = ref.frame_orientation_cost(q, frame, R_des.flatten(order="F"), W, 1, eps) + fd = _central(lambda qq: ref.frame_orientation_cost(qq, frame, R_des.flatten(order="F"), W, 1, eps)[0], q) + assert np.allclose(grad[:nv], fd, atol=1e-6), np.max(np.abs(grad[:nv] - fd)) + assert np.count_nonzero(grad[nv:]) == 0 and np.count_nonzero(hess[nv:, :]) == 0 + # the Gauss-Newton structure J_e^T diag(a'') J_e + Jw = np.asarray(ref.frame_jacobian(q, frame, "LOCAL"))[3:, :] + Je = ref._so3_right_jacobian_inv(phi) @ Jw + curv = ref._activation(1, W, phi, eps)[2] + assert np.allclose(hess[:nv, :nv], Je.T @ (curv[:, None] * Je)) + # the log map at pi/2 pitch (the Euler singularity) is regular + Rz = lambda a: np.array([[np.cos(a), -np.sin(a), 0], [np.sin(a), np.cos(a), 0], [0, 0, 1]]) + Ry = lambda a: np.array([[np.cos(a), 0, np.sin(a)], [0, 1, 0], [-np.sin(a), 0, np.cos(a)]]) + R = Rz(0.3) @ Ry(np.pi / 2) + phi = ref._so3_log(R) + from scipy.spatial.transform import Rotation + assert np.allclose(Rotation.from_rotvec(phi).as_matrix(), R, atol=1e-10) + # near pi branch + R_pi = Rotation.from_rotvec(np.array([0.0, 0.0, np.pi - 1e-7])).as_matrix() + assert np.isclose(np.linalg.norm(ref._so3_log(R_pi)), np.pi - 1e-7, atol=1e-6) + + +@pytest.mark.parametrize("name", ["iiwa14", "fr3"]) +def test_frame_velocity_cost(name): + ref = _ref(name, False) + nq, nv = ref.robot.get_num_pos(), ref.robot.get_num_vel() + rng = np.random.default_rng(12) + q = rng.standard_normal(nq); qd = rng.standard_normal(nv) + x = np.concatenate([q, qd]) + W = np.arange(1, 7, dtype=float); v_des = np.array([0.1, 0.0, -0.1, 0.2, 0.0, 0.1]); eps = 0.05 + for frame in _frames(ref): + val, grad, hess = ref.frame_velocity_cost(q, qd, frame, v_des, W, 1, eps) + fd = _central(lambda xx: ref.frame_velocity_cost(xx[:nq], xx[nq:], frame, v_des, W, 1, eps)[0], x) + assert np.allclose(grad, fd, atol=1e-6), np.max(np.abs(grad - fd)) + # the configuration derivative of J qd against a finite difference of the Jacobian itself + A = ref._frame_velocity_q_derivative(q, qd, frame) + F = np.zeros((6, nv)) + for k in range(nq): + qp = q.copy(); qp[k] += _H; qm = q.copy(); qm[k] -= _H + F[:, k] = (ref.frame_jacobian(qp, frame) @ qd - ref.frame_jacobian(qm, frame) @ qd) / (2 * _H) + assert np.allclose(A, F, atol=1e-6) + J = np.asarray(ref.frame_jacobian(q, frame)) + Jx = np.hstack([A, J]); curv = ref._activation(1, W, J @ qd - v_des, eps)[2] + assert np.allclose(hess, Jx.T @ (curv[:, None] * Jx)) + go2 = _ref("go2", True) + with pytest.raises(NotImplementedError): + go2.frame_velocity_cost(go2.robot.normalize_floating_base_q_input(rng.standard_normal(go2.robot.get_num_pos())), + rng.standard_normal(go2.robot.get_num_vel()), _frames(go2)[0], v_des, W, 1, eps) + + +def test_barrier_activations(): + W = np.array([0.1, 0.2, 0.3]); vals = np.array([0.5, 1.95, -0.3]) + lo = np.array([0.0, 0.0, -np.inf]); hi = np.array([2.0, 2.0, 0.0]); delta = 0.1 + for kind in (2, 3): + val, sl, cu = RBDReference._activation(kind, W, vals, delta, vals, lo, hi) + fd = _central(lambda vv: np.sum(RBDReference._activation(kind, W, vv, delta, vv, lo, hi)[0]), vals) + assert np.allclose(sl, fd, atol=1e-6) + fd2 = _central(lambda vv: np.sum(RBDReference._activation(kind, W, vv, delta, vv, lo, hi)[1]), vals) + assert np.allclose(cu, fd2, atol=1e-4) + # relaxed log barrier: log above delta, finite quadratic extension below and beyond the bound + v, d1, d2 = RBDReference._relaxed_log_barrier_1d(np.array([0.5, 0.1, 0.05, -1.0]), 0.1, delta) + assert np.isclose(v[0], -0.1 * np.log(0.5)) and np.isclose(v[1], -0.1 * np.log(0.1)) + assert np.isfinite(v).all() and v[3] > v[2] > v[1] + # continuity of value and slope at delta + a = RBDReference._relaxed_log_barrier_1d(np.array([delta + 1e-9]), 0.1, delta) + b = RBDReference._relaxed_log_barrier_1d(np.array([delta - 1e-9]), 0.1, delta) + assert abs(a[0][0] - b[0][0]) < 1e-7 and abs(a[1][0] - b[1][0]) < 1e-6 + # quadratic barrier is zero inside and the infinite bound contributes nothing + val, sl, cu = RBDReference._activation(3, W, vals, 0.0, vals, lo, hi) + assert val[0] == 0.0 and sl[0] == 0.0 and cu[0] == 0.0 and val[2] == 0.0 + assert np.isclose(val[2], 0.0) and np.isclose(RBDReference._activation(3, W, np.array([-0.3]), 0.0, np.array([-0.3]), np.array([-np.inf]), np.array([-0.5]))[0][0], 0.5 * 0.1 * 0.04) + with pytest.raises(ValueError): + RBDReference._activation(2, W, vals, delta) diff --git a/tests/test_joint_dynamics_friction_model.py b/tests/test_joint_dynamics_friction_model.py new file mode 100644 index 0000000..2c40ada --- /dev/null +++ b/tests/test_joint_dynamics_friction_model.py @@ -0,0 +1,166 @@ +"""Joint-dynamics friction models: sign (default) and tanh (smooth). + +RBDReference(use_joint_dynamics=True, friction_model=, friction_width=) applies +tau += damping*qd + friction*s(qd) with s = sign under the default model and +s = tanh(qd/w) under the smooth model. The smooth model has derivatives at and +near qd = 0, which is where the sign model has none: the inverse-dynamics +gradient gains (friction/w) sech^2(qd/w) on the dqd diagonal and the +second-order qd-qd block gains -(2 friction/w^2) tanh(qd/w) sech^2(qd/w) on +its triple diagonal. Every check here is analytic against the reference's own +finite differences, at qd = 0 and at |qd| <= w, on fr3 (damping, friction and +a mimic joint, so the per-slot fold is exercised) and on iiwa14 floating +(damping only, root skipped). +""" +from pathlib import Path + +import numpy as np +import pytest + +from URDFParser import URDFParser +from RBDReference import RBDReference + +_ASSETS = Path(__file__).resolve().parents[1] / "robot_assets" +_W = 0.03 + + +def _robot(name, floating=False): + return URDFParser().parse(str(_ASSETS / f"{name}.urdf"), floating_base=floating) + + +def _samples(rng, nv, width): + """qd at rest, inside the width, at the width, and well outside it.""" + at_rest = np.zeros(nv) + inside = rng.uniform(-0.5, 0.5, nv) * width + edge = np.sign(rng.standard_normal(nv)) * width + outside = np.sign(rng.standard_normal(nv)) * (0.5 + np.abs(rng.standard_normal(nv))) + return [at_rest, inside, edge, outside] + + +def _central(f, x, k, h): + xp = x.copy(); xp[k] += h + xm = x.copy(); xm[k] -= h + return (np.asarray(f(xp)) - np.asarray(f(xm))) / (2.0 * h) + + +def test_invalid_model_and_width_rejected(): + robot = _robot("fr3") + with pytest.raises(ValueError): + RBDReference(robot, use_joint_dynamics=True, friction_model="smooth") + with pytest.raises(ValueError): + RBDReference(robot, use_joint_dynamics=True, friction_model="tanh", friction_width=0.0) + with pytest.raises(ValueError): + RBDReference(robot, use_joint_dynamics=True, friction_model="tanh", + friction_width=np.ones(robot.get_num_vel() + 1)) + + +def test_sign_model_is_the_default_and_unchanged(): + robot = _robot("fr3") + ref = RBDReference(robot, use_joint_dynamics=True) + assert ref.friction_model == "sign" + rng = np.random.default_rng(1) + nv = robot.get_num_vel() + q = rng.standard_normal(nv) + for qd in _samples(rng, nv, _W): + bias = ref._joint_dynamics_bias(qd) + expected = np.zeros(nv) + for jid in range(robot.get_num_bodies()): + k = robot.get_joint_index_v(jid) + alpha = ref._mimic_multiplier(jid) + expected[k] += alpha * (robot.get_damping_by_id(jid) * qd[k] + + robot.get_friction_by_id(jid) * np.sign(qd[k])) + assert np.allclose(bias, expected, rtol=0, atol=1e-12) + # the sign model carries no friction slope anywhere and no second derivative + assert np.allclose(ref._joint_dynamics_derivative_diag(qd, 2), 0.0) + grad = ref._joint_dynamics_derivative_diag(qd, 1) + damping_only = np.zeros(nv) + for jid in range(robot.get_num_bodies()): + k = robot.get_joint_index_v(jid) + damping_only[k] += ref._mimic_multiplier(jid) * robot.get_damping_by_id(jid) + assert np.allclose(grad, damping_only, rtol=0, atol=1e-12) + + +@pytest.mark.parametrize("width", [_W, "per_slot"], ids=["scalar", "per_slot"]) +def test_tanh_bias_and_slopes_match_closed_form(width): + robot = _robot("fr3") + nv = robot.get_num_vel() + rng = np.random.default_rng(2) + w = np.full(nv, _W) if width == _W else rng.uniform(0.01, 0.1, nv) + ref = RBDReference(robot, use_joint_dynamics=True, friction_model="tanh", + friction_width=(_W if width == _W else w)) + for qd in _samples(rng, nv, w): + bias = np.zeros(nv); g1 = np.zeros(nv); g2 = np.zeros(nv) + for jid in range(robot.get_num_bodies()): + k = robot.get_joint_index_v(jid) + a = ref._mimic_multiplier(jid) + b = robot.get_damping_by_id(jid) + c = robot.get_friction_by_id(jid) + th = np.tanh(qd[k] / w[k]); s2 = 1.0 - th * th + bias[k] += a * (b * qd[k] + c * th) + g1[k] += a * (b + c / w[k] * s2) + g2[k] += a * (-2.0 * c / w[k] ** 2 * th * s2) + assert np.allclose(ref._joint_dynamics_bias(qd), bias, rtol=0, atol=1e-12) + assert np.allclose(ref._joint_dynamics_derivative_diag(qd, 1), g1, rtol=0, atol=1e-12) + assert np.allclose(ref._joint_dynamics_derivative_diag(qd, 2), g2, rtol=0, atol=1e-12) + # the friction slope is present at rest: this is the point the sign model misses + assert np.max(ref._joint_dynamics_derivative_diag(np.zeros(nv), 1)) > 1.0 + + +@pytest.mark.parametrize("name,floating", [("fr3", False), ("iiwa14", True)], + ids=["fr3-fixed", "iiwa14-floating"]) +def test_tanh_gradient_matches_finite_difference(name, floating): + robot = _robot(name, floating) + ref = RBDReference(robot, use_joint_dynamics=True, friction_model="tanh", friction_width=_W) + nv = robot.get_num_vel() + rng = np.random.default_rng(3) + q = ref.robot.normalize_floating_base_q_input(rng.standard_normal(robot.get_num_pos())) \ + if floating else rng.standard_normal(nv) + qdd = rng.standard_normal(nv) + h = 1e-6 + for qd in _samples(rng, nv, _W): + grad = np.asarray(ref.inverse_dynamics_gradient(q, qd, qdd, GRAVITY=-9.81)) + dqd_analytic = grad[:, nv:] + fd = np.zeros((nv, nv)) + for k in range(nv): + fd[:, k] = _central(lambda v: ref.inverse_dynamics(q, v, qdd, GRAVITY=-9.81)[0], qd, k, h) + err = np.max(np.abs(dqd_analytic - fd)) + assert err < 1e-5, f"{name} qd={qd}: dc_dqd analytic vs FD {err:.2e}" + # forward-dynamics gradient follows by composition + u = rng.standard_normal(nv) + _, qdd_dqd = ref.forward_dynamics_gradient(q, qd, u) + fd2 = np.zeros((nv, nv)) + for k in range(nv): + fd2[:, k] = _central(lambda v: ref.forward_dynamics(q, v, u), qd, k, h) + err2 = np.max(np.abs(qdd_dqd - fd2)) / max(1.0, np.max(np.abs(fd2))) + assert err2 < 1e-5, f"{name} qd={qd}: qdd_dqd analytic vs FD {err2:.2e}" + + +@pytest.mark.parametrize("name,floating", [("fr3", False), ("iiwa14", True)], + ids=["fr3-fixed", "iiwa14-floating"]) +def test_tanh_second_order_matches_finite_difference_of_gradient(name, floating): + robot = _robot(name, floating) + ref = RBDReference(robot, use_joint_dynamics=True, friction_model="tanh", friction_width=_W) + nv = robot.get_num_vel() + rng = np.random.default_rng(4) + q = ref.robot.normalize_floating_base_q_input(rng.standard_normal(robot.get_num_pos())) \ + if floating else rng.standard_normal(nv) + qdd = rng.standard_normal(nv) + h = 1e-5 + for qd in _samples(rng, nv, _W): + _, d2tau_dqd, _, _ = ref.idsva_so(q, qd, qdd, GRAVITY=-9.81) + # d2tau[i, j, k] = d/dqd_k of dc_dqd[i, j] + for k in range(nv): + fd = _central(lambda v: np.asarray(ref.inverse_dynamics_gradient(q, v, qdd, GRAVITY=-9.81))[:, nv:], qd, k, h) + err = np.max(np.abs(d2tau_dqd[:, :, k] - fd)) / max(1.0, np.max(np.abs(fd))) + assert err < 1e-4, f"{name} qd={qd} k={k}: d2tau_dqd vs FD of gradient {err:.2e}" + if name == "fr3" and np.all(np.abs(qd) <= _W): + # the friction term is the dominant entry on the triple diagonal at |qd| <= w + so = ref._joint_dynamics_derivative_diag(qd, 2) + if np.any(qd != 0.0): + assert np.max(np.abs(so)) > 1.0 + # fdsva_so follows by composition: its qd-qd block against FD of the FD gradient + u = rng.standard_normal(nv) + _, _, daba_dvdv, _ = ref.fdsva_so(q, qd, u, GRAVITY=-9.81) + for k in range(nv): + fd = _central(lambda v: ref.forward_dynamics_gradient(q, v, u)[1], qd, k, h) + err = np.max(np.abs(daba_dvdv[:, :, k] - fd)) / max(1.0, np.max(np.abs(fd))) + assert err < 1e-4, f"{name} qd={qd} k={k}: daba_dvdv vs FD {err:.2e}"