phyz 0.3.0

Multi-physics differentiable simulation engine
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104

//! Recursive Newton-Euler Algorithm (RNEA) — inverse dynamics.
//!
//! Given (q, v, qdd), compute the required joint torques tau.

use crate::math::{DVec, SpatialTransform, SpatialVec, Vec3};
use crate::model::{JointType, Model, State};

/// Compute joint velocity contribution S * qd for any joint type.
fn joint_velocity(joint: &crate::model::Joint, qd: &[f64]) -> SpatialVec {
    match joint.joint_type {
        JointType::Revolute | JointType::Hinge => SpatialVec::new(joint.axis * qd[0], Vec3::zero()),
        JointType::Prismatic | JointType::Slide => {
            SpatialVec::new(Vec3::zero(), joint.axis * qd[0])
        }
        JointType::Spherical | JointType::Ball => {
            SpatialVec::new(Vec3::new(qd[0], qd[1], qd[2]), Vec3::zero())
        }
        JointType::Free => SpatialVec::new(
            Vec3::new(qd[0], qd[1], qd[2]),
            Vec3::new(qd[3], qd[4], qd[5]),
        ),
        JointType::Fixed => SpatialVec::zero(),
    }
}

/// Run RNEA: compute inverse dynamics torques.
///
/// Given state (q, v) and desired accelerations `qdd`, returns the torques needed.
pub fn rnea(model: &Model, state: &State, qdd: &DVec) -> DVec {
    let nb = model.nbodies();
    let mut tau = DVec::zeros(model.nv);

    let mut x_tree = vec![SpatialTransform::identity(); nb];
    let mut vel = vec![SpatialVec::zero(); nb];
    let mut acc = vec![SpatialVec::zero(); nb];

    let a0 = SpatialVec::new(Vec3::zero(), -model.gravity);

    // ── Forward pass: velocities and accelerations ──
    for i in 0..nb {
        let body = &model.bodies[i];
        let joint = &model.joints[body.joint_idx];
        let q_idx = model.q_offsets[body.joint_idx];
        let v_idx = model.v_offsets[body.joint_idx];
        let ndof = joint.ndof();

        let x_joint = if ndof == 0 {
            SpatialTransform::identity()
        } else {
            let q_slice = &state.q.as_slice()[q_idx..q_idx + ndof];
            joint.joint_transform_slice(q_slice)
        };
        x_tree[i] = x_joint.compose(&joint.parent_to_joint);

        let v_joint = joint_velocity(joint, &state.v.as_slice()[v_idx..v_idx + ndof]);
        let a_joint = joint_velocity(joint, &qdd.as_slice()[v_idx..v_idx + ndof]);

        if body.parent < 0 {
            vel[i] = v_joint;
            acc[i] = x_tree[i].apply_motion(&a0) + a_joint;
        } else {
            let pi = body.parent as usize;
            let v_parent = x_tree[i].apply_motion(&vel[pi]);
            vel[i] = v_parent + v_joint;

            let a_parent = x_tree[i].apply_motion(&acc[pi]);
            acc[i] = a_parent + vel[i].cross_motion(&v_joint) + a_joint;
        }
    }

    // ── Backward pass: forces and torques ──
    let mut forces = vec![SpatialVec::zero(); nb];
    for i in 0..nb {
        let body = &model.bodies[i];
        let ia = body.inertia.to_matrix();
        forces[i] = ia.mul_vec(&acc[i]) + vel[i].cross_force(&ia.mul_vec(&vel[i]));
    }

    for i in (0..nb).rev() {
        let body = &model.bodies[i];
        let joint = &model.joints[body.joint_idx];
        let v_idx = model.v_offsets[body.joint_idx];
        let ndof = joint.ndof();

        if ndof == 1 {
            let s_i = joint.motion_subspace();
            tau[v_idx] = s_i.dot(&forces[i]);
        } else if ndof > 1 {
            let s_mat = joint.motion_subspace_matrix(); // 6 x ndof
            let f_vec = crate::aba::sv_to_dvec(&forces[i]);
            let phyz_vec = s_mat.transpose().mul_vec(&f_vec); // ndof x 1
            for k in 0..ndof {
                tau[v_idx + k] = phyz_vec[k];
            }
        }

        if body.parent >= 0 {
            let pi = body.parent as usize;
            forces[pi] = forces[pi] + x_tree[i].inv_apply_force(&forces[i]);
        }
    }

    tau
}