feat(python): add more control other multibodies throug the python bindings

bindings/python/rapier-py-3d/python/rapier3d/_rapier3d.pyi

@@ -758,6 +758,10 @@ class MultibodyLink:
     def coords(self) -> list[float]: ...
     @property
     def joint_rot(self) -> Rotation3: ...
+    @property
+    def assembly_id(self) -> int: ...
+    @property
+    def ndofs(self) -> int: ...
 
 class Multibody:
     @property
@@ -772,6 +776,16 @@ class Multibody:
     def set_self_contacts_enabled(self, v: bool) -> None: ...
     def damping(self) -> list[float]: ...
     def set_damping(self, values: list[float]) -> None: ...
+    def armature(self) -> list[float]: ...
+    def set_armature(self, values: list[float]) -> None: ...
+    def generalized_position(self) -> list[float]: ...
+    def set_link_motor(
+        self, link_id: int, axis: JointAxis, target_pos: float,
+        target_vel: float, stiffness: float, damping: float,
+    ) -> None: ...
+    def set_link_motor_max_force(self, link_id: int, axis: JointAxis, max_force: float) -> None: ...
+    def set_link_motor_model(self, link_id: int, axis: JointAxis, model: MotorModel) -> None: ...
+    def link_motor(self, link_id: int, axis: JointAxis) -> JointMotor | None: ...
     def generalized_velocity(self) -> list[float]: ...
     def set_generalized_velocity(self, values: list[float]) -> None: ...
     def apply_displacements(self, displacements: Sequence[float]) -> None: ...
@@ -1052,7 +1066,21 @@ class MjcfActuatorHandle:
     @property
     def name(self) -> str | None: ...
     @property
+    def kind(self) -> str: ...
+    @property
+    def joint_name(self) -> str | None: ...
+    @property
     def joint(self) -> ImpulseJointHandle | MultibodyJointHandle | None: ...
+    @property
+    def gear(self) -> list[float]: ...
+    @property
+    def ctrl_range(self) -> list[float] | None: ...
+    @property
+    def force_range(self) -> list[float] | None: ...
+    @property
+    def kp(self) -> float | None: ...
+    @property
+    def kv(self) -> float | None: ...
 
 class MjcfContactHooks:
     def filter_contact_pair(self, ctx: PairFilterContext) -> SolverFlags | None: ...
@@ -1068,6 +1096,14 @@ class MjcfRobotHandles:
     @property
     def actuators(self) -> list[MjcfActuatorHandle]: ...
     @property
+    def body_names(self) -> list[str | None]: ...
+    @property
+    def joint_names(self) -> list[str | None]: ...
+    @property
+    def body_name_to_idx(self) -> dict[str, int]: ...
+    @property
+    def joint_name_to_idx(self) -> dict[str, int]: ...
+    @property
     def keyframe_names(self) -> list[str | None]: ...
     def apply_controls(
         self,
@@ -1096,6 +1132,14 @@ class MjcfRobot:
     def gravity(self) -> Vec3: ...
     @property
     def keyframe_names(self) -> list[str | None]: ...
+    @property
+    def body_names(self) -> list[str | None]: ...
+    @property
+    def joint_names(self) -> list[str | None]: ...
+    @property
+    def body_name_to_idx(self) -> dict[str, int]: ...
+    @property
+    def joint_name_to_idx(self) -> dict[str, int]: ...
     def append_transform(self, transform: IsometryLike) -> None: ...
     def insert_using_impulse_joints(
         self, bodies: RigidBodySet, colliders: ColliderSet, impulse_joints: ImpulseJointSet

bindings/python/rapier-py-3d/src/joints.rs

@@ -3265,6 +3265,29 @@ impl MultibodyLink {
     fn joint_rot(&self) -> Rotation3 {
         Rotation3(self.0.joint().joint_rot().into())
     }
+    /// Index of this link's first degree of freedom in the
+    /// articulation-wide generalized vectors.
+    ///
+    /// The link owns the :py:attr:`ndofs` entries starting at this
+    /// index in :py:meth:`Multibody.generalized_position`,
+    /// :py:meth:`Multibody.generalized_velocity`, and friends.
+    #[getter]
+    fn assembly_id(&self) -> usize {
+        self.0.assembly_id()
+    }
+    /// Number of degrees of freedom of this link's joint.
+    #[getter]
+    fn ndofs(&self) -> usize {
+        self.0.joint().ndofs()
+    }
+}
+
+/// Indices, in the 6-entry spatial layout of
+/// `MultibodyJoint::coords`, of the joint's unlocked axes, in the
+/// order the solver assigns its degrees of freedom.
+fn free_dof_axes(joint: &rapier::dynamics::MultibodyJoint) -> impl Iterator<Item = usize> {
+    let locked = joint.data.locked_axes.bits();
+    (0..rapier::math::SPATIAL_DIM).filter(move |i| locked & (1u8 << *i) == 0)
 }
 
 // =================================================================
@@ -3330,6 +3353,27 @@ impl Multibody {
             Ok(f(mb))
         })
     }
+    /// Run `f` on link `link_id`, or raise `IndexError` if there is no
+    /// such link.
+    fn with_link_mut(
+        &mut self,
+        link_id: usize,
+        f: impl FnOnce(&mut rapier::dynamics::MultibodyLink),
+    ) -> PyResult<()> {
+        self.with_mut(|mb| match mb.link_mut(link_id) {
+            Some(link) => {
+                f(link);
+                Ok(())
+            }
+            None => Err(link_index_error(link_id)),
+        })?
+    }
+}
+
+/// The `IndexError` raised when a link id does not name a link of the
+/// articulation.
+fn link_index_error(link_id: usize) -> PyErr {
+    crate::pyo3::exceptions::PyIndexError::new_err(format!("no multibody link with id {link_id}"))
 }
 
 #[pymethods]
@@ -3395,6 +3439,119 @@ impl Multibody {
             Ok(())
         })?
     }
+    /// Per-degree-of-freedom armature (reflected rotor inertia,
+    /// length :attr:`ndofs`).
+    ///
+    /// Armature is added straight to the generalized mass-matrix
+    /// diagonal, which is what makes a stiff position motor stable at
+    /// large gains.
+    fn armature(&self) -> PyResult<Vec<Real>> {
+        self.with_ref(|mb| mb.armature().iter().copied().collect())
+    }
+    /// Set the per-DOF armature values.
+    ///
+    /// :raises ValueError: If ``values`` length differs from `ndofs`.
+    fn set_armature(&mut self, values: Vec<Real>) -> PyResult<()> {
+        self.with_mut(|mb| {
+            let a = mb.armature_mut();
+            if values.len() != a.len() {
+                return Err(PyValueError::new_err(format!(
+                    "expected {} armature values (ndofs), got {}",
+                    a.len(),
+                    values.len()
+                )));
+            }
+            for (i, v) in values.iter().enumerate() {
+                a[i] = *v;
+            }
+            Ok(())
+        })?
+    }
+    /// Generalized position vector (one entry per DOF).
+    ///
+    /// Entry ``link.assembly_id + k`` is the ``k``-th unlocked
+    /// coordinate of that link's joint, so the layout matches
+    /// :py:meth:`generalized_velocity` and the vectors accepted by
+    /// :py:meth:`apply_displacements`.
+    fn generalized_position(&self) -> PyResult<Vec<Real>> {
+        self.with_ref(|mb| {
+            let mut out = vec![0.0; mb.ndofs()];
+            for link in mb.links() {
+                let joint = link.joint();
+                let coords = joint.coords();
+                for (k, axis) in free_dof_axes(joint).enumerate() {
+                    out[link.assembly_id() + k] = coords[axis];
+                }
+            }
+            out
+        })
+    }
+    /// Fully configure the motor on ``axis`` of link ``link_id``'s
+    /// joint (the reduced-coordinate counterpart of
+    /// :py:meth:`GenericJoint.set_motor`).
+    ///
+    /// :raises IndexError: If ``link_id`` is out of range.
+    fn set_link_motor(
+        &mut self,
+        link_id: usize,
+        axis: JointAxis,
+        target_pos: Real,
+        target_vel: Real,
+        stiffness: Real,
+        damping: Real,
+    ) -> PyResult<()> {
+        self.with_link_mut(link_id, |link| {
+            link.joint
+                .data
+                .set_motor(axis.to_rapier(), target_pos, target_vel, stiffness, damping);
+        })
+    }
+    /// Clamp the maximum force the motor on ``axis`` of link
+    /// ``link_id`` can apply.
+    ///
+    /// :raises IndexError: If ``link_id`` is out of range.
+    fn set_link_motor_max_force(
+        &mut self,
+        link_id: usize,
+        axis: JointAxis,
+        max_force: Real,
+    ) -> PyResult<()> {
+        self.with_link_mut(link_id, |link| {
+            link.joint
+                .data
+                .set_motor_max_force(axis.to_rapier(), max_force);
+        })
+    }
+    /// Select the motor model on ``axis`` of link ``link_id`` (see
+    /// :class:`MotorModel`).
+    ///
+    /// :raises IndexError: If ``link_id`` is out of range.
+    fn set_link_motor_model(
+        &mut self,
+        link_id: usize,
+        axis: JointAxis,
+        model: MotorModel,
+    ) -> PyResult<()> {
+        self.with_link_mut(link_id, |link| {
+            link.joint
+                .data
+                .set_motor_model(axis.to_rapier(), model.to_rapier());
+        })
+    }
+    /// Return the motor configured on ``axis`` of link ``link_id``,
+    /// if any.
+    ///
+    /// :raises IndexError: If ``link_id`` is out of range.
+    fn link_motor(&self, link_id: usize, axis: JointAxis) -> PyResult<Option<JointMotor>> {
+        self.with_ref(|mb| {
+            let link = mb.link(link_id).ok_or_else(|| link_index_error(link_id))?;
+            Ok(link
+                .joint()
+                .data
+                .motor(axis.to_rapier())
+                .map(JointMotor::from_rapier))
+        })?
+    }
     /// Generalized velocity vector (one entry per DOF).
     fn generalized_velocity(&self) -> PyResult<Vec<Real>> {
         self.with_ref(|mb| mb.generalized_velocity().iter().copied().collect())