[software] 添加16DOF早期训练仿真与Sim2Real闭环

05_software/train/rc_mjlab/mujoco_sim/dynamics.py

@@ -0,0 +1,97 @@
+"""Pinocchio dynamics: FK, Jacobian, IK for the wheeled-legged robot."""
+
+import numpy as np
+import pinocchio as pin
+from config import MJCF_PATH, LEG_NAMES
+
+# Foot frame names in Pinocchio model (wheel link centers)
+FOOT_FRAMES = {leg: f"{leg}_wheel_Link" for leg in LEG_NAMES}
+
+# Leg joint names for each leg
+_LEG_JOINT_NAMES = {
+    leg: [f"{leg}_{jt}" for jt in ("hip_abduction_joint", "hip_pitch_joint", "knee_joint")]
+    for leg in LEG_NAMES
+}
+
+
+class Dynamics:
+    """Pinocchio-based kinematics/dynamics. Deployable on real hardware."""
+
+    def __init__(self):
+        self.model = pin.buildModelFromMJCF(str(MJCF_PATH))
+        self.data = self.model.createData()
+
+        # Cache frame IDs
+        self._foot_fids = {}
+        for leg, fname in FOOT_FRAMES.items():
+            self._foot_fids[leg] = self.model.getFrameId(fname)
+
+        # Cache joint velocity indices for each leg (3 joints)
+        self._leg_v_indices = {}
+        for leg, jnames in _LEG_JOINT_NAMES.items():
+            indices = []
+            for jn in jnames:
+                jid = self.model.getJointId(jn)
+                indices.append(self.model.joints[jid].idx_v)
+            self._leg_v_indices[leg] = indices
+
+        # Cache joint config indices for each leg
+        self._leg_q_indices = {}
+        for leg, jnames in _LEG_JOINT_NAMES.items():
+            indices = []
+            for jn in jnames:
+                jid = self.model.getJointId(jn)
+                indices.append(self.model.joints[jid].idx_q)
+            self._leg_q_indices[leg] = indices
+
+    def update(self, q: np.ndarray, dq: np.ndarray):
+        """Forward kinematics + Jacobians.
+
+        Args:
+            q: Pinocchio config (nq=23: pos3, quat_xyzw4, joints16)
+            dq: Pinocchio velocity (nv=22: v_body3, w_body3, joints16)
+        """
+        pin.forwardKinematics(self.model, self.data, q, dq)
+        pin.updateFramePlacements(self.model, self.data)
+        pin.computeJointJacobians(self.model, self.data, q)
+
+    def get_foot_pos(self, leg: str) -> np.ndarray:
+        """Foot (wheel center) position in world frame (3,)."""
+        return self.data.oMf[self._foot_fids[leg]].translation.copy()
+
+    def get_foot_jacobian_leg(self, leg: str) -> np.ndarray:
+        """3x3 linear Jacobian of foot w.r.t. 3 leg joints (world frame)."""
+        fid = self._foot_fids[leg]
+        J_full = pin.getFrameJacobian(
+            self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
+        cols = self._leg_v_indices[leg]
+        return J_full[:, cols]
+
+    def inverse_kinematics(self, leg: str, target_pos: np.ndarray,
+                           q_current: np.ndarray, max_iter=30, eps=1e-4) -> np.ndarray:
+        """Numerical IK for one leg. Returns (3,) joint angles.
+
+        Args:
+            leg: Leg name
+            target_pos: Desired foot position in world frame (3,)
+            q_current: Current full Pinocchio config (nq=23)
+        """
+        q = q_current.copy()
+        fid = self._foot_fids[leg]
+        q_indices = self._leg_q_indices[leg]
+
+        for _ in range(max_iter):
+            pin.forwardKinematics(self.model, self.data, q)
+            pin.updateFramePlacements(self.model, self.data)
+            err = target_pos - self.data.oMf[fid].translation
+            if np.linalg.norm(err) < eps:
+                break
+            pin.computeJointJacobians(self.model, self.data, q)
+            J = pin.getFrameJacobian(
+                self.model, self.data, fid, pin.LOCAL_WORLD_ALIGNED)[:3, :]
+            J_leg = J[:, self._leg_v_indices[leg]]
+            dq = np.linalg.solve(J_leg.T @ J_leg + 1e-6 * np.eye(3), J_leg.T @ err)
+            for i, idx in enumerate(q_indices):
+                q[idx] += dq[i]
+
+        return np.array([q[idx] for idx in q_indices])