mujoco_rs/wrappers/
mj_data.rs

1use super::mj_model::{MjModel, MjtSameFrame, MjtObj};
2use std::io::{self, Error, ErrorKind};
3use super::mj_auxiliary::MjContact;
4use super::mj_primitive::*;
5use crate::mujoco_c::*;
6use std::ffi::CString;
7use std::ptr;
8
9use crate::{mj_view_indices, mj_model_nx_to_mapping, mj_model_nx_to_nitem};
10use crate::{view_creator, fixed_size_info_method, info_with_view};
11
12/*******************************************/
13// Types
14
15
16pub type MjtState = mjtState;
17/*******************************************/
18
19
20/**************************************************************************************************/
21// MjData
22/**************************************************************************************************/
23
24/// Wrapper around the ``mjData`` struct.
25/// Provides lifetime guarantees as well as automatic cleanup.
26pub struct MjData<'a> {
27    data: *mut mjData,
28    model: &'a MjModel
29}
30
31// Allow usage in threaded contexts as the data won't be shared anywhere outside Rust,
32// except in the C++ code.
33unsafe impl Send for MjData<'_> {}
34unsafe impl Sync for MjData<'_> {}
35
36
37impl<'a> MjData<'a> {
38    /// Constructor for a new MjData. This should is called from MjModel.
39    pub fn new(model: &'a MjModel) -> Self {
40        unsafe {
41            Self {
42                data: mj_makeData(model.ffi()),
43                model: model,
44            }
45        }
46    }
47
48    /// Reference to the wrapped FFI struct.
49    pub fn ffi(&self) -> &mjData {
50        unsafe { self.data.as_ref().unwrap() }
51    }
52
53    /// Mutable reference to the wrapped FFI struct.
54    pub unsafe fn ffi_mut(&mut self) -> &mut mjData {
55        unsafe { self.data.as_mut().unwrap() }
56    }
57
58    /// Returns a slice of detected contacts.
59    /// To obtain the contact force, call [`MjData::contact_force`].
60    pub fn contacts(&self) -> &[MjContact] {
61        unsafe {
62            std::slice::from_raw_parts((*self.data).contact, (*self.data).ncon as usize)
63        }
64    }
65
66    /// Obtains a [`MjActuatorDataInfo`] struct containing information about the name, id, and
67    /// indices required for obtaining a slice view to the correct locations in [`MjData`].
68    /// The actual view can be obtained via [`MjActuatorDataInfo::view`].
69    pub fn actuator(&self, name: &str) -> Option<MjActuatorDataInfo> {
70        let c_name = CString::new(name).unwrap();
71        let id = unsafe { mj_name2id(self.model.ffi(), mjtObj::mjOBJ_ACTUATOR as i32, c_name.as_ptr())};
72        if id == -1 {  // not found
73            return None;
74        }
75
76        let ctrl;
77        let act;
78        let model_ffi = self.model.ffi();
79        unsafe {
80            ctrl = (id as usize, 1);
81            act = mj_view_indices!(id as usize, model_ffi.actuator_actadr, model_ffi.nu as usize, model_ffi.na as usize);
82        }
83
84        Some(MjActuatorDataInfo { name: name.to_string(), id: id as usize, ctrl, act})
85    }
86
87    fixed_size_info_method! { Data, model.ffi(), body, [xfrc_applied: 6, xpos: 3, xquat: 4, xmat: 9, xipos: 3, ximat: 9, subtree_com: 3, cinert: 10, crb: 10, cvel: 6, subtree_linvel: 3, subtree_angmom: 3, cacc: 6, cfrc_int: 6, cfrc_ext: 6] }
88    fixed_size_info_method! { Data, model.ffi(), camera, [xpos: 3, xmat: 9] }
89    fixed_size_info_method! { Data, model.ffi(), geom, [xpos: 3, xmat: 9] }
90    fixed_size_info_method! { Data, model.ffi(), site, [xpos: 3, xmat: 9] }
91    fixed_size_info_method! { Data, model.ffi(), light, [xpos: 3, xdir: 3] }
92
93
94    /// Obtains a [`MjJointDataInfo`] struct containing information about the name, id, and
95    /// indices required for obtaining a slice view to the correct locations in [`MjData`].
96    /// The actual view can be obtained via [`MjJointDataInfo::view`].
97    pub fn joint(&self, name: &str) -> Option<MjJointDataInfo> {
98        let c_name = CString::new(name).unwrap();
99        let id = unsafe { mj_name2id(self.model.ffi(), mjtObj::mjOBJ_JOINT as i32, c_name.as_ptr())};
100        if id == -1 {  // not found
101            return None;
102        }
103        let model_ffi = self.model.ffi();
104        let id = id as usize;
105        unsafe {
106            let nq_range = mj_view_indices!(id, mj_model_nx_to_mapping!(model_ffi, nq), mj_model_nx_to_nitem!(model_ffi, nq), model_ffi.nq);
107            let nv_range = mj_view_indices!(id, mj_model_nx_to_mapping!(model_ffi, nv), mj_model_nx_to_nitem!(model_ffi, nv), model_ffi.nv);
108
109            // $id:expr, $addr_map:expr, $njnt:expr, $max_n:expr
110            let qpos = nq_range;
111            let qvel = nv_range;
112            let qacc_warmstart = nv_range;
113            let qfrc_applied = nv_range;
114            let qacc = nv_range;
115            let xanchor = (id * 3, 3);
116            let xaxis = (id * 3, 3);
117            #[allow(non_snake_case)]
118            let qLDiagInv = nv_range;
119            let qfrc_bias = nv_range;
120            let qfrc_passive = nv_range;
121            let qfrc_actuator = nv_range;
122            let qfrc_smooth = nv_range;
123            let qacc_smooth = nv_range;
124            let qfrc_constraint = nv_range;
125            let qfrc_inverse = nv_range;
126            
127            /* Special case attributes, used for some internal calculation */
128            // cdof
129            // cdof_dot
130
131            Some(MjJointDataInfo {name: name.to_string(), id: id as usize,
132                qpos, qvel, qacc_warmstart, qfrc_applied, qacc, xanchor, xaxis, qLDiagInv, qfrc_bias,
133                qfrc_passive, qfrc_actuator, qfrc_smooth, qacc_smooth, qfrc_constraint, qfrc_inverse
134            })
135        }
136    }
137
138    /// Obtains a [`MjSensorDataInfo`] struct containing information about the name, id, and
139    /// indices required for obtaining a slice view to the correct locations in [`MjData`].
140    /// The actual view can be obtained via [`MjJointDataInfo::view`].
141    pub fn sensor(&self, name: &str) -> Option<MjSensorDataInfo> {
142        let c_name = CString::new(name).unwrap();
143        let id = unsafe { mj_name2id(self.model.ffi(), mjtObj::mjOBJ_SENSOR as i32, c_name.as_ptr())};
144        if id == -1 {  // not found
145            return None;
146        }
147        let model_ffi = self.model.ffi();
148        let id = id as usize;
149
150        unsafe {
151            let data = mj_view_indices!(id, mj_model_nx_to_mapping!(model_ffi, nsensordata), mj_model_nx_to_nitem!(model_ffi, nsensordata), model_ffi.nsensordata);
152            Some(MjSensorDataInfo { id, name: name.to_string(), data })
153        }
154    }
155
156
157    /// Obtains a [`MjTendonDataInfo`] struct containing information about the name, id, and
158    /// indices required for obtaining a slice view to the correct locations in [`MjData`].
159    /// The actual view can be obtained via [`MjJointDataInfo::view`].
160    #[allow(non_snake_case)]
161    pub fn tendon(&self, name: &str) -> Option<MjTendonDataInfo> {
162        let c_name = CString::new(name).unwrap();
163        let id = unsafe { mj_name2id(self.model.ffi(), mjtObj::mjOBJ_TENDON as i32, c_name.as_ptr())};
164        if id == -1 {  // not found
165            return None;
166        }
167
168        let model_ffi = self.model.ffi();
169        let id = id as usize;
170        let nv = model_ffi.nv as usize;
171        let wrapadr = (id, 1);
172        let wrapnum = (id, 1);
173        let J_rownnz = (id, 1);
174        let J_rowadr = (id, 1);
175        let J_colind = (id * nv, nv);
176        let length = (id, 1);
177        let J = (id * nv, nv);
178        let velocity = (id, 1);
179
180        Some(MjTendonDataInfo { id, name: name.to_string(), wrapadr, wrapnum, J_rownnz, J_rowadr, J_colind, length, J, velocity })
181    }
182
183    /// Steps the MuJoCo simulation.
184    pub fn step(&mut self) {
185        unsafe {
186            mj_step(self.model.ffi(), self.ffi_mut());
187        }
188    }
189
190    /// Calculates new dynamics. This is a wrapper around `mj_step1`.
191    pub fn step1(&mut self) {
192        unsafe {
193            mj_step1(self.model.ffi(), self.ffi_mut());
194        }
195    }
196
197    /// Calculates the rest after dynamics and integrates in time.
198    /// This is a wrapper around `mj_step2`.
199    pub fn step2(&mut self) {
200        unsafe {
201            mj_step2(self.model.ffi(), self.ffi_mut());
202        }
203    }
204
205    /// Forward dynamics: same as mj_step but do not integrate in time.
206    /// This is a wrapper around `mj_forward`.
207    pub fn forward(&mut self) {
208        unsafe { 
209            mj_forward(self.model.ffi(), self.ffi_mut());
210        }
211    }
212
213    /// [`MjData::forward`] dynamics with skip.
214    /// This is a wrapper around `mj_forwardSkip`.
215    pub fn forward_skip(&mut self, skipstage: mjtStage, skipsensor: bool) {
216        unsafe { 
217            mj_forwardSkip(self.model.ffi(), self.ffi_mut(), skipstage as i32, skipsensor as i32);
218        }
219    }
220
221    /// Inverse dynamics: qacc must be set before calling ([`MjData::forward`]).
222    /// This is a wrapper around `mj_inverse`.
223    pub fn inverse(&mut self) {
224        unsafe {
225            mj_inverse(self.model.ffi(), self.ffi_mut());
226        }
227    }
228
229    /// [`MjData::inverse`] dynamics with skip; skipstage is mjtStage.
230    /// This is a wrapper around `mj_inverseSkip`.
231    pub fn inverse_skip(&mut self, skipstage: mjtStage, skipsensor: bool) {
232        unsafe {
233            mj_inverseSkip(self.model.ffi(), self.ffi_mut(), skipstage as i32, skipsensor as i32);
234        }
235    }
236
237    /// Calculates the contact force for the given `contact_id`.
238    /// The `contact_id` matches the index of the contact when iterating
239    /// via [`MjData::contacts`].
240    /// Calls `mj_contactForce` internally.
241    pub fn contact_force(&self, contact_id: usize) -> [f64; 6] {
242        let mut force = [0.0; 6];
243        unsafe {
244            mj_contactForce(
245                self.model.ffi(), self.data,
246                contact_id as i32, force.as_mut_ptr()
247            );
248        }
249        force
250    }
251
252    /* Partially auto-generated */
253
254    /// Reset data to defaults.
255    pub fn reset(&mut self) {
256        unsafe { mj_resetData(self.model.ffi(), self.ffi_mut()) }
257    }
258
259    /// Reset data to defaults, fill everything else with debug_value.
260    pub fn reset_debug(&mut self, debug_value: u8) {
261        unsafe { mj_resetDataDebug(self.model.ffi(), self.ffi_mut(), debug_value) }
262    }
263
264    /// Reset data. If 0 <= key < nkey, set fields from specified keyframe.
265    pub fn reset_keyframe(&mut self, key: i32) {
266        unsafe { mj_resetDataKeyframe(self.model.ffi(), self.ffi_mut(), key) }
267    }
268
269    /// Print mjData to text file, specifying format.
270    /// float_format must be a valid printf-style format string for a single float value
271    pub fn print_formatted(&self, filename: &str, float_format: &str) {
272        let c_filename = CString::new(filename).unwrap();
273        let c_float_format = CString::new(float_format).unwrap();
274        unsafe { mj_printFormattedData(self.model.ffi(), self.ffi(), c_filename.as_ptr(), c_float_format.as_ptr()) }
275    }
276
277    /// Print data to text file.
278    pub fn print(&self, filename: &str) {
279        let c_filename = CString::new(filename).unwrap();
280        unsafe { mj_printData(self.model.ffi(), self.ffi(), c_filename.as_ptr()) }
281    }
282
283    /// Run position-dependent computations.
284    pub fn fwd_position(&mut self) {
285        unsafe { mj_fwdPosition(self.model.ffi(), self.ffi_mut()) }
286    }
287
288    /// Run velocity-dependent computations.
289    pub fn fwd_velocity(&mut self) {
290        unsafe { mj_fwdVelocity(self.model.ffi(), self.ffi_mut()) }
291    }
292
293    /// Compute actuator force qfrc_actuator.
294    pub fn fwd_actuation(&mut self) {
295        unsafe { mj_fwdActuation(self.model.ffi(), self.ffi_mut()) }
296    }
297
298    /// Add up all non-constraint forces, compute qacc_smooth.
299    pub fn fwd_acceleration(&mut self) {
300        unsafe { mj_fwdAcceleration(self.model.ffi(), self.ffi_mut()) }
301    }
302
303    /// Run selected constraint solver.
304    pub fn fwd_constraint(&mut self) {
305        unsafe { mj_fwdConstraint(self.model.ffi(), self.ffi_mut()) }
306    }
307
308    /// Euler integrator, semi-implicit in velocity.
309    pub fn euler(&mut self) {
310        unsafe { mj_Euler(self.model.ffi(), self.ffi_mut()) }
311    }
312
313    /// Runge-Kutta explicit order-N integrator.
314    pub fn runge_kutta(&mut self, n: i32) {
315        unsafe { mj_RungeKutta(self.model.ffi(), self.ffi_mut(), n) }
316    }
317
318    /// Implicit-in-velocity integrators.
319    pub fn implicit(&mut self) {
320        unsafe { mj_implicit(self.model.ffi(), self.ffi_mut()) }
321    }
322
323    /// Run position-dependent computations in inverse dynamics.
324    pub fn inv_position(&mut self) {
325        unsafe { mj_invPosition(self.model.ffi(), self.ffi_mut()) }
326    }
327
328    /// Run velocity-dependent computations in inverse dynamics.
329    pub fn inv_velocity(&mut self) {
330        unsafe { mj_invVelocity(self.model.ffi(), self.ffi_mut()) }
331    }
332
333    /// Apply the analytical formula for inverse constraint dynamics.
334    pub fn inv_constraint(&mut self) {
335        unsafe { mj_invConstraint(self.model.ffi(), self.ffi_mut()) }
336    }
337
338    /// Compare forward and inverse dynamics, save results in fwdinv.
339    pub fn compare_fwd_inv(&mut self) {
340        unsafe { mj_compareFwdInv(self.model.ffi(), self.ffi_mut()) }
341    }
342
343    /// Evaluate position-dependent sensors.
344    pub fn sensor_pos(&mut self) {
345        unsafe { mj_sensorPos(self.model.ffi(), self.ffi_mut()) }
346    }
347
348    /// Evaluate velocity-dependent sensors.
349    pub fn sensor_vel(&mut self) {
350        unsafe { mj_sensorVel(self.model.ffi(), self.ffi_mut()) }
351    }
352
353    /// Evaluate acceleration and force-dependent sensors.
354    pub fn sensor_acc(&mut self) {
355        unsafe { mj_sensorAcc(self.model.ffi(), self.ffi_mut()) }
356    }
357
358    /// Evaluate position-dependent energy (potential).
359    pub fn energy_pos(&mut self) {
360        unsafe { mj_energyPos(self.model.ffi(), self.ffi_mut()) }
361    }
362
363    /// Evaluate velocity-dependent energy (kinetic).
364    pub fn energy_vel(&mut self) {
365        unsafe { mj_energyVel(self.model.ffi(), self.ffi_mut()) }
366    }
367
368    /// Check qpos, reset if any element is too big or nan.
369    pub fn check_pos(&mut self) {
370        unsafe { mj_checkPos(self.model.ffi(), self.ffi_mut()) }
371    }
372
373    /// Check qvel, reset if any element is too big or nan.
374    pub fn check_vel(&mut self) {
375        unsafe { mj_checkVel(self.model.ffi(), self.ffi_mut()) }
376    }
377
378    /// Check qacc, reset if any element is too big or nan.
379    pub fn check_acc(&mut self) {
380        unsafe { mj_checkAcc(self.model.ffi(), self.ffi_mut()) }
381    }
382
383    /// Run forward kinematics.
384    pub fn kinematics(&mut self) {
385        unsafe { mj_kinematics(self.model.ffi(), self.ffi_mut()) }
386    }
387
388    /// Map inertias and motion dofs to global frame centered at CoM.
389    pub fn com_pos(&mut self) {
390        unsafe { mj_comPos(self.model.ffi(), self.ffi_mut()) }
391    }
392
393    /// Compute camera and light positions and orientations.
394    pub fn camlight(&mut self) {
395        unsafe { mj_camlight(self.model.ffi(), self.ffi_mut()) }
396    }
397
398    /// Compute flex-related quantities.
399    pub fn flex_comp(&mut self) {
400        unsafe { mj_flex(self.model.ffi(), self.ffi_mut()) }
401    }
402
403    /// Compute tendon lengths, velocities and moment arms.
404    pub fn tendon_comp(&mut self) {
405        unsafe { mj_tendon(self.model.ffi(), self.ffi_mut()) }
406    }
407
408    /// Compute actuator transmission lengths and moments.
409    pub fn transmission(&mut self) {
410        unsafe { mj_transmission(self.model.ffi(), self.ffi_mut()) }
411    }
412
413    /// Run composite rigid body inertia algorithm (CRB).
414    pub fn crb(&mut self) {
415        unsafe { mj_crb(self.model.ffi(), self.ffi_mut()) }
416    }
417
418    /// Make inertia matrix.
419    pub fn make_m(&mut self) {
420        unsafe { mj_makeM(self.model.ffi(), self.ffi_mut()) }
421    }
422
423    /// Compute sparse L'*D*L factorizaton of inertia matrix.
424    pub fn factor_m(&mut self) {
425        unsafe { mj_factorM(self.model.ffi(), self.ffi_mut()) }
426    }
427
428    /// Compute cvel, cdof_dot.
429    pub fn com_vel(&mut self) {
430        unsafe { mj_comVel(self.model.ffi(), self.ffi_mut()) }
431    }
432
433    /// Compute qfrc_passive from spring-dampers, gravity compensation and fluid forces.
434    pub fn passive(&mut self) {
435        unsafe { mj_passive(self.model.ffi(), self.ffi_mut()) }
436    }
437
438    /// Sub-tree linear velocity and angular momentum: compute subtree_linvel, subtree_angmom.
439    pub fn subtree_vel(&mut self) {
440        unsafe { mj_subtreeVel(self.model.ffi(), self.ffi_mut()) }
441    }
442
443    /// RNE: compute M(qpos)*qacc + C(qpos,qvel); flg_acc=false removes inertial term.
444    pub fn rne(&mut self, flg_acc: bool) -> Vec<MjtNum> {
445        let mut out = vec![0.0; self.model.ffi().nv as usize];
446        unsafe { mj_rne(self.model.ffi(), self.ffi_mut(), flg_acc as i32, out.as_mut_ptr()) };
447        out
448    }
449
450    /// RNE with complete data: compute cacc, cfrc_ext, cfrc_int.
451    pub fn rne_post_constraint(&mut self) {
452        unsafe { mj_rnePostConstraint(self.model.ffi(), self.ffi_mut()) }
453    }
454
455    /// Run collision detection.
456    pub fn collision(&mut self) {
457        unsafe { mj_collision(self.model.ffi(), self.ffi_mut()) }
458    }
459
460    /// Construct constraints.
461    pub fn make_constraint(&mut self) {
462        unsafe { mj_makeConstraint(self.model.ffi(), self.ffi_mut()) }
463    }
464
465    /// Find constraint islands.
466    pub fn island(&mut self) {
467        unsafe { mj_island(self.model.ffi(), self.ffi_mut()) }
468    }
469
470    /// Compute inverse constraint inertia efc_AR.
471    pub fn project_constraint(&mut self) {
472        unsafe { mj_projectConstraint(self.model.ffi(), self.ffi_mut()) }
473    }
474
475    /// Compute efc_vel, efc_aref.
476    pub fn reference_constraint(&mut self) {
477        unsafe { mj_referenceConstraint(self.model.ffi(), self.ffi_mut()) }
478    }
479
480    /// Compute efc_state, efc_force, qfrc_constraint, and (optionally) cone Hessians.
481    /// If cost is not NULL, set *cost = s(jar) where jar = Jac*qacc-aref.
482    /// Nullable: cost
483    pub fn constraint_update(&mut self, jar: &[MjtNum], cost: Option<&mut MjtNum>, flg_cone_hessian: bool) {
484        unsafe { mj_constraintUpdate(
485            self.model.ffi(), self.ffi_mut(),
486            jar.as_ptr(), cost.map_or(ptr::null_mut(), |x| x as *mut MjtNum),
487            flg_cone_hessian as i32
488        ) }
489    }
490
491    /// Add contact to d->contact list; return 0 if success; 1 if buffer full.
492    pub fn add_contact(&mut self, con: &MjContact) -> io::Result<()> {
493        match unsafe { mj_addContact(self.model.ffi(), self.ffi_mut(), con) } {
494            0 => Ok(()),
495            1 => Err(Error::new(ErrorKind::StorageFull, "buffer full")),
496            _ => Err(Error::new(ErrorKind::Other, "unknown error"))
497        }
498    }
499
500    /// Compute 3/6-by-nv end-effector Jacobian of a global point attached to the given body.
501    /// Set `jacp` to `true` to calculate the translational Jacobian and `jacr` to `true` for
502    /// the rotational Jacobian. Returns a `(Vec, Vec)` for translation and rotation. Empty `Vec`s
503    /// indicate that the corresponding Jacobian was not computed.
504    pub fn jac(&self, jacp: bool, jacr: bool, point: &[MjtNum; 3], body_id: i32) -> (Vec<MjtNum>, Vec<MjtNum>) {
505        let required_len = 3 * self.model.ffi().nv as usize;
506        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
507        let mut jacr_vec = if jacr { vec![0 as MjtNum; required_len] } else { vec![] };
508
509        unsafe {
510            mj_jac(
511                self.model.ffi(),
512                self.ffi(),
513                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
514                if jacr { jacr_vec.as_mut_ptr() } else { ptr::null_mut() },
515                point.as_ptr(),
516                body_id,
517            )
518        };
519
520        (jacp_vec, jacr_vec)
521    }
522
523    /// Compute body frame end-effector Jacobian.
524    /// Set `jacp`/`jacr` to `true` to calculate translational/rotational components.
525    /// Returns `(Vec, Vec)` for translation and rotation. Empty `Vec`s indicate not computed.
526    pub fn jac_body(&self, jacp: bool, jacr: bool, body_id: i32) -> (Vec<MjtNum>, Vec<MjtNum>) {
527        let required_len = 3 * self.model.ffi().nv as usize;
528        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
529        let mut jacr_vec = if jacr { vec![0 as MjtNum; required_len] } else { vec![] };
530
531        unsafe {
532            mj_jacBody(
533                self.model.ffi(),
534                self.ffi(),
535                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
536                if jacr { jacr_vec.as_mut_ptr() } else { ptr::null_mut() },
537                body_id,
538            )
539        };
540
541        (jacp_vec, jacr_vec)
542    }
543
544    /// Compute body center-of-mass end-effector Jacobian.
545    /// Set `jacp`/`jacr` to `true` to calculate translational/rotational components.
546    /// Returns `(Vec, Vec)` for translation and rotation. Empty `Vec`s indicate not computed.
547    pub fn jac_body_com(&self, jacp: bool, jacr: bool, body_id: i32) -> (Vec<MjtNum>, Vec<MjtNum>) {
548        let required_len = 3 * self.model.ffi().nv as usize;
549        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
550        let mut jacr_vec = if jacr { vec![0 as MjtNum; required_len] } else { vec![] };
551
552        unsafe {
553            mj_jacBodyCom(
554                self.model.ffi(),
555                self.ffi(),
556                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
557                if jacr { jacr_vec.as_mut_ptr() } else { ptr::null_mut() },
558                body_id,
559            )
560        };
561
562        (jacp_vec, jacr_vec)
563    }
564
565    /// Compute subtree center-of-mass end-effector Jacobian (translational only).
566    /// Set `jacp` to `true` to calculate the translational component. Returns a `Vec`.
567    /// Empty `Vec` indicates that the Jacobian was not computed.
568    pub fn jac_subtree_com(&mut self, jacp: bool, body_id: i32) -> Vec<MjtNum> {
569        let required_len = 3 * self.model.ffi().nv as usize;
570        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
571
572        unsafe {
573            mj_jacSubtreeCom(
574                self.model.ffi(),
575                self.ffi_mut(),
576                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
577                body_id,
578            )
579        };
580
581        jacp_vec
582    }
583
584    /// Compute geom end-effector Jacobian.
585    /// Set `jacp`/`jacr` to `true` to calculate translational/rotational components.
586    /// Returns `(Vec, Vec)` for translation and rotation. Empty `Vec`s indicate not computed.
587    pub fn jac_geom(&self, jacp: bool, jacr: bool, geom_id: i32) -> (Vec<MjtNum>, Vec<MjtNum>) {
588        let required_len = 3 * self.model.ffi().nv as usize;
589        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
590        let mut jacr_vec = if jacr { vec![0 as MjtNum; required_len] } else { vec![] };
591
592        unsafe {
593            mj_jacGeom(
594                self.model.ffi(),
595                self.ffi(),
596                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
597                if jacr { jacr_vec.as_mut_ptr() } else { ptr::null_mut() },
598                geom_id,
599            )
600        };
601
602        (jacp_vec, jacr_vec)
603    }
604
605    /// Compute site end-effector Jacobian.
606    /// Set `jacp`/`jacr` to `true` to calculate translational/rotational components.
607    /// Returns `(Vec, Vec)` for translation and rotation. Empty `Vec`s indicate not computed.
608    pub fn jac_site(&self, jacp: bool, jacr: bool, site_id: i32) -> (Vec<MjtNum>, Vec<MjtNum>) {
609        let required_len = 3 * self.model.ffi().nv as usize;
610        let mut jacp_vec = if jacp { vec![0 as MjtNum; required_len] } else { vec![] };
611        let mut jacr_vec = if jacr { vec![0 as MjtNum; required_len] } else { vec![] };
612
613        unsafe {
614            mj_jacSite(
615                self.model.ffi(),
616                self.ffi(),
617                if jacp { jacp_vec.as_mut_ptr() } else { ptr::null_mut() },
618                if jacr { jacr_vec.as_mut_ptr() } else { ptr::null_mut() },
619                site_id,
620            )
621        };
622
623        (jacp_vec, jacr_vec)
624    }
625
626    /// Compute subtree angular momentum matrix.
627    pub fn angmom_mat(&mut self, body_id: i32) -> Vec<MjtNum> {
628        let mut mat = vec![0.0; 3 * self.model.ffi().nv as usize];
629        unsafe { mj_angmomMat(self.model.ffi(), self.ffi_mut(), mat.as_mut_ptr(), body_id) };
630        mat
631    }
632
633    /// Compute object 6D velocity (rot:lin) in object-centered frame, world/local orientation.
634    pub fn object_velocity(&self, obj_type: MjtObj, obj_id: i32, flg_local: bool) -> [MjtNum; 6] {
635        let mut result: [MjtNum; 6] = [0.0; 6];
636        unsafe { mj_objectVelocity(
637            self.model.ffi(), self.ffi(),
638            obj_type as i32, obj_id,
639            result.as_mut_ptr(), flg_local as i32
640        ) };
641        result
642    }
643
644    /// Compute object 6D acceleration (rot:lin) in object-centered frame, world/local orientation.
645    pub fn object_acceleration(&self, obj_type: MjtObj, obj_id: i32, flg_local: bool) -> [MjtNum; 6] {
646        let mut result: [MjtNum; 6] = [0.0; 6];
647        unsafe { mj_objectAcceleration(
648            self.model.ffi(), self.ffi(),
649            obj_type as i32, obj_id,
650            result.as_mut_ptr(), flg_local as i32
651        ) };
652        result
653    }
654
655    /// Returns smallest signed distance between two geoms and optionally segment from geom1 to geom2.
656    pub fn geom_distance(&self, geom1_id: i32, geom2_id: i32, dist_max: MjtNum, fromto: Option<&mut [MjtNum; 6]>) -> MjtNum {
657        unsafe { mj_geomDistance(
658            self.model.ffi(), self.ffi(),
659            geom1_id, geom2_id, dist_max,
660            fromto.map_or(ptr::null_mut(), |x| x.as_mut_ptr())
661        ) }
662    }
663
664    /// Map from body local to global Cartesian coordinates, sameframe takes values from mjtSameFrame.
665    /// Returns (global position, global orientation matrix).
666    /// Wraps ``mj_local2Global``.
667    pub fn local_to_global(&mut self, pos: &[MjtNum; 3], quat: &[MjtNum; 4], body_id: i32, sameframe: MjtSameFrame) -> ([MjtNum; 3], [MjtNum; 9]) {
668        /* Create uninitialized because this gets filled by the function. */
669        let mut xpos: [MjtNum; 3] =  [0.0; 3];
670        let mut xmat: [MjtNum; 9] = [0.0; 9];
671        unsafe { mj_local2Global(self.ffi_mut(), xpos.as_mut_ptr(), xmat.as_mut_ptr(), pos.as_ptr(), quat.as_ptr(), body_id, sameframe as MjtByte) };
672        (xpos, xmat)
673    }
674
675    /// Intersect multiple rays emanating from a single point.
676    /// Similar semantics to mj_ray, but vec is an array of (nray x 3) directions.
677    /// Returns (geomids, distances).
678    pub fn multi_ray(
679        &mut self, pnt: &[MjtNum; 3], vec: &[[MjtNum; 3]], geomgroup: Option<&[MjtByte; mjNGROUP as usize]>,
680        flg_static: bool, bodyexclude: i32, cutoff: MjtNum
681    ) -> (Vec<i32>, Vec<MjtNum>) {
682        let nray = vec.len();
683        let mut geom_id = vec![0; nray];
684        let mut distance = vec![0.0; nray];
685
686        unsafe { mj_multiRay(
687            self.model.ffi(), self.ffi_mut(), pnt.as_ptr(),
688            vec.as_ptr() as *const MjtNum, geomgroup.map_or(ptr::null(), |x| x.as_ptr()),
689            flg_static as u8, bodyexclude, geom_id.as_mut_ptr(),
690            distance.as_mut_ptr(), nray as i32, cutoff
691        ) };
692
693        (geom_id, distance)
694    }
695
696    /// Intersect ray (pnt+x*vec, x>=0) with visible geoms, except geoms in bodyexclude.
697    /// Return distance (x) to nearest surface, or -1 if no intersection and output geomid.
698    /// geomgroup, flg_static are as in mjvOption; geomgroup==NULL skips group exclusion.
699    pub fn ray(
700        &self, pnt: &[MjtNum; 3], vec: &[MjtNum; 3],
701        geomgroup: Option<&[MjtByte; mjNGROUP as usize]>, flg_static: bool, bodyexclude: i32
702    ) -> (i32, MjtNum) {
703        let mut geom_id = -1;
704        let dist = unsafe { mj_ray(
705            self.model.ffi(), self.ffi(),
706            pnt.as_ptr(), vec.as_ptr(),
707            geomgroup.map_or(ptr::null(), |x| x.as_ptr()),
708            flg_static as MjtByte, bodyexclude, &mut geom_id
709        ) };
710        (geom_id, dist)
711    }
712
713    /// Returns a direct pointer to the underlying model.
714    /// THIS IS NOT TO BE USED.
715    /// It is only meant for the viewer code, which currently still depends
716    /// on mutable pointers to model and data. This will be removed in the future.
717    pub(crate) unsafe fn __raw(&self) -> *mut mjData {
718        self.data
719    }
720
721}
722
723impl Drop for MjData<'_> {
724    fn drop(&mut self) {
725        unsafe {
726            mj_deleteData(self.data);
727        }
728    }
729}
730
731
732/**************************************************************************************************/
733// Joint view
734/**************************************************************************************************/
735info_with_view!(
736    Data,
737    joint,
738    [
739        qpos: f64, qvel: f64, qacc_warmstart: f64, qfrc_applied: f64, qacc: f64, xanchor: f64, xaxis: f64, qLDiagInv: f64, qfrc_bias: f64,
740        qfrc_passive: f64, qfrc_actuator: f64, qfrc_smooth: f64, qacc_smooth: f64, qfrc_constraint: f64, qfrc_inverse: f64
741    ],
742    []
743);
744
745/// Deprecated name for [`MjJointDataInfo`].
746#[deprecated]
747pub type MjJointInfo = MjJointDataInfo;
748
749
750/* Backward compatibility */
751impl MjJointDataViewMut<'_> {
752    /// Deprecated. Use [`MjJointDataViewMut::zero`] instead.
753    #[deprecated]
754    pub fn reset(&mut self) {
755        self.zero();
756    }
757}
758
759/// Deprecated name for [`MjJointDataView`].
760#[deprecated]
761pub type MjJointView<'d> = MjJointDataView<'d>;
762
763/// Deprecated name for [`MjJointDataViewMut`].
764#[deprecated]
765pub type MjJointViewMut<'d> = MjJointDataViewMut<'d>;
766
767/**************************************************************************************************/
768// Sensor view
769/**************************************************************************************************/
770info_with_view!(Data, sensor, sensor, [data: f64], []);
771
772/**************************************************************************************************/
773// Geom view
774/**************************************************************************************************/
775info_with_view!(Data, geom, geom_, [xpos: f64, xmat: f64], []);
776
777/// Deprecated name for [`MjGeomDataInfo`].
778#[deprecated]
779pub type MjGeomInfo = MjGeomDataInfo;
780
781/// Deprecated name for [`MjGeomDataView`].
782#[deprecated]
783pub type MjGeomView<'d> = MjGeomDataView<'d>;
784
785/// Deprecated name for [`MjGeomDataViewMut`].
786#[deprecated]
787pub type MjGeomViewMut<'d> = MjGeomDataViewMut<'d>;
788
789/**************************************************************************************************/
790// Actuator view
791/**************************************************************************************************/
792info_with_view!(Data, actuator, [ctrl: f64], [act: f64]);
793
794/// Deprecated name for [`MjActuatorDataInfo`].
795#[deprecated]
796pub type MjActuatorInfo = MjActuatorDataInfo;
797
798/// Deprecated name for [`MjActuatorDataView`].
799#[deprecated]
800pub type MjActuatorView<'d> = MjActuatorDataView<'d>;
801
802/// Deprecated name for [`MjActuatorDataViewMut`].
803#[deprecated]
804pub type MjActuatorViewMut<'d> = MjActuatorDataViewMut<'d>;
805
806/**************************************************************************************************/
807// Body view
808/**************************************************************************************************/
809info_with_view!(
810    Data, body, [
811        xfrc_applied: f64, xpos: f64, xquat: f64, xmat: f64, xipos: f64, ximat: f64,
812        subtree_com: f64, cinert: f64, crb: f64, cvel: f64, subtree_linvel: f64,
813        subtree_angmom: f64, cacc: f64, cfrc_int: f64, cfrc_ext: f64
814    ], []
815);
816
817/**************************************************************************************************/
818// Camera view
819/**************************************************************************************************/
820info_with_view!(Data, camera, cam_, [xpos: f64, xmat: f64], []);
821
822/**************************************************************************************************/
823// Site view
824/**************************************************************************************************/
825info_with_view!(Data, site, site_, [xpos: f64, xmat: f64], []);
826
827/**************************************************************************************************/
828// Tendon view
829/**************************************************************************************************/
830info_with_view!(Data, tendon, ten_, [wrapadr: i32, wrapnum: i32, J_rownnz: i32, J_rowadr: i32, J_colind: i32, length: f64, J: f64, velocity: f64], []);
831
832/**************************************************************************************************/
833// Light view
834/**************************************************************************************************/
835info_with_view!(Data, light, light_, [xpos: f64, xdir: f64], []);
836
837/**************************************************************************************************/
838// Unit tests
839/**************************************************************************************************/
840
841#[cfg(test)]
842mod test {
843    use crate::prelude::*;
844    use super::*;
845
846    const MODEL: &str = "
847<mujoco>
848  <worldbody>
849    <light ambient=\"0.2 0.2 0.2\"/>
850    <body name=\"ball\">
851        <geom name=\"green_sphere\" pos=\".2 .2 .1\" size=\".1\" rgba=\"0 1 0 1\" solref=\"0.004 1.0\"/>
852        <joint name=\"ball_joint\" type=\"free\"/>
853    </body>
854
855    <geom name=\"floor1\" type=\"plane\" size=\"10 10 1\" solref=\"0.004 1.0\"/>
856  </worldbody>
857</mujoco>
858";
859
860    macro_rules! assert_almost_eq {
861        ($a:expr, $b:expr) => {
862            assert!(($a - $b).abs() < 1e-6)
863        }
864    }
865
866    #[test]
867    fn test_body_view() {
868        let model = MjModel::from_xml_string(MODEL).unwrap();
869        let mut data = model.make_data();
870        let body_info = data.body("ball").unwrap();
871        let mut cvel;
872
873        data.step1();
874
875        for _ in 0..10 {
876            data.step2();
877            data.step1();  // step() and step2() update before integration, thus we need to manually update non-state variables.
878        }
879
880        // The ball should start in a still position
881        cvel = body_info.view(&data).cvel;
882        assert_almost_eq!(cvel[0], 0.0);
883        assert_almost_eq!(cvel[1], 0.0);
884        assert_almost_eq!(cvel[2], 0.0);
885        assert_almost_eq!(cvel[3], 0.0);
886        assert_almost_eq!(cvel[4], 0.0);
887        // assert_almost_eq!(cvel[5], 0.0);  // Ignore due to slight instability of the model.
888
889        // Give the ball some velocity
890        body_info.view_mut(&mut data).xfrc_applied[0] = 5.0;
891        data.step2();
892        data.step1();
893
894        let view = body_info.view(&data);
895        cvel = view.cvel;
896        println!("{:?}", cvel);
897        assert_almost_eq!(cvel[0], 0.0);
898        assert!(cvel[1] > 0.0);  // wy should be positive when rolling with positive vx.
899        assert_almost_eq!(cvel[2], 0.0);
900        assert!(cvel[3] > 0.0);  // vx should point in the direction of the applied force.
901        assert_almost_eq!(cvel[4], 0.0);  // vy should be 0.
902        // assert_almost_eq!(cvel[5], 0.0);  // vz should be 0, but we don't test it due to jumpiness (instability) of the ball.
903
904        assert_almost_eq!(view.xfrc_applied[0], 5.0); // the original force should stay applied.
905
906        data.step2();
907        data.step1();
908    }
909
910    #[test]
911    fn test_copy_reset_variants() {
912        let model = MjModel::from_xml_string(MODEL).unwrap();
913        let mut data = model.make_data();
914
915        // Test reset variants
916        data.reset();
917        data.reset_debug(7);
918        data.reset_keyframe(0);
919    }
920
921    #[test]
922    fn test_dynamics_and_sensors() {
923        let model = MjModel::from_xml_string(MODEL).unwrap();
924        let mut data = model.make_data();
925
926        // Simulation pipeline components
927        data.fwd_position();
928        data.fwd_velocity();
929        data.fwd_actuation();
930        data.fwd_acceleration();
931        data.fwd_constraint();
932
933        data.euler();
934        data.runge_kutta(4);
935        // data.implicit();  // integrator isn't implicit in the model => skip this check
936
937        data.inv_position();
938        data.inv_velocity();
939        data.inv_constraint();
940        data.compare_fwd_inv();
941
942        // Sensors
943        data.sensor_pos();
944        data.sensor_vel();
945        data.sensor_acc();
946
947        data.energy_pos();
948        data.energy_vel();
949
950        data.check_pos();
951        data.check_vel();
952        data.check_acc();
953
954        data.kinematics();
955        data.com_pos();
956        data.camlight();
957        data.flex_comp();
958        data.tendon_comp();
959        data.transmission();
960        data.crb();
961        data.make_m();
962        data.factor_m();
963        data.com_vel();
964        data.passive();
965        data.subtree_vel();
966    }
967
968
969    #[test]
970    fn test_rne_and_collision_pipeline() {
971        let model = MjModel::from_xml_string(MODEL).unwrap();
972        let mut data = model.make_data();
973        data.step();
974
975        // mj_rne returns a scalar as result
976        data.rne(true);
977
978        data.rne_post_constraint();
979
980        // // Collision and constraint pipeline
981        data.collision();
982        data.make_constraint();
983        data.island();
984        data.project_constraint();
985        data.reference_constraint();
986
987        let jar = vec![0.0; (data.model.ffi().nv) as usize];
988        let mut cost = 0.0;
989        data.constraint_update(&jar, None, false);
990        data.constraint_update(&jar, Some(&mut cost), true);
991    }
992
993    #[test]
994    fn test_add_contact() {
995        let model = MjModel::from_xml_string(MODEL).unwrap();
996        let mut data = model.make_data();
997
998        // Add a dummy contact
999        let dummy_contact = unsafe { std::mem::zeroed() };
1000        data.add_contact(&dummy_contact).unwrap();
1001    }
1002
1003    #[test]
1004    fn test_jacobian() {
1005        let model = MjModel::from_xml_string(MODEL).unwrap();
1006        let mut data = model.make_data();
1007
1008        let nv = data.model.ffi().nv as usize;
1009        let expected_len = 3 * nv;
1010
1011        // Use a small offset point relative to the joint origin
1012        let point = [0.1, 0.0, 0.0];
1013
1014        let ball_body_id = model.body("ball").unwrap().id as i32;
1015
1016        // Test global point Jacobian
1017        let (jacp, jacr) = data.jac(true, true, &point, ball_body_id);
1018        assert_eq!(jacp.len(), expected_len);
1019        assert_eq!(jacr.len(), expected_len);
1020
1021        // Test body frame Jacobian
1022        let (jacp_body, jacr_body) = data.jac_body(true, true, ball_body_id);
1023        assert_eq!(jacp_body.len(), expected_len);
1024        assert_eq!(jacr_body.len(), expected_len);
1025
1026        // Test body COM Jacobian
1027        let (jacp_com, jacr_com) = data.jac_body_com(true, true, ball_body_id);
1028        assert_eq!(jacp_com.len(), expected_len);
1029        assert_eq!(jacr_com.len(), expected_len);
1030
1031        // Test subtree COM Jacobian (translational only)
1032        let jac_subtree = data.jac_subtree_com(true, 0);
1033        assert_eq!(jac_subtree.len(), expected_len);
1034
1035        // Test geom Jacobian
1036        let (jacp_geom, jacr_geom) = data.jac_geom(true, true, ball_body_id);
1037        assert_eq!(jacp_geom.len(), expected_len);
1038        assert_eq!(jacr_geom.len(), expected_len);
1039
1040        // Test site Jacobian
1041        let (jacp_site, jacr_site) = data.jac_site(true, true, ball_body_id);
1042        assert_eq!(jacp_site.len(), expected_len);
1043        assert_eq!(jacr_site.len(), expected_len);
1044
1045        // Test flags set to false produce empty Vec
1046        let (jacp_none, jacr_none) = data.jac(false, false, &[0.0; 3], ball_body_id);
1047        assert!(jacp_none.is_empty());
1048        assert!(jacr_none.is_empty());
1049    }
1050
1051    #[test]
1052    fn test_angmom_and_object_dynamics() {
1053        let model = MjModel::from_xml_string(MODEL).unwrap();
1054        let mut data = model.make_data();
1055
1056        let mat = data.angmom_mat(0);
1057        assert_eq!(mat.len(), (3 * data.model.ffi().nv as usize));
1058
1059        let vel = data.object_velocity(MjtObj::mjOBJ_BODY, 0, true);
1060        assert_eq!(vel, vel); // just ensure it returns 6-length
1061
1062        let acc = data.object_acceleration(MjtObj::mjOBJ_BODY, 0, false);
1063        assert_eq!(acc, acc);
1064    }
1065
1066    #[test]
1067    fn test_geom_distance_and_transforms() {
1068        let model = MjModel::from_xml_string(MODEL).unwrap();
1069        let mut data = model.make_data();
1070        data.step();
1071
1072        let mut ft = [0.0; 6];
1073        let dist = data.geom_distance(0, 0, 1.0, Some(&mut ft));
1074        assert_eq!(ft, [0.0; 6]);
1075        assert_eq!(dist, 1.0);
1076
1077        let pos = [0.0; 3];
1078        let quat = [1.0, 0.0, 0.0, 0.0];
1079        let (xpos, xmat) = data.local_to_global(&pos, &quat, 0, MjtSameFrame::mjSAMEFRAME_NONE);
1080        assert_eq!(xpos.len(), 3);
1081        assert_eq!(xmat.len(), 9);
1082
1083        let ray_vecs = [[1.0, 0.0, 0.0], [1.0, 0.0, 0.0], [1.0, 0.0, 0.0]];
1084        let rays = data.multi_ray(&pos, &ray_vecs, None, false, -1, 10.0);
1085        assert_eq!(rays.0.len(), 3);
1086        assert_eq!(rays.1.len(), 3);
1087
1088        let (geomid, dist) = data.ray(&pos, &[1.0, 0.0, 0.0], None, true, -1);
1089        assert!(dist.is_finite());
1090        assert!(geomid >= -1);
1091    }
1092}
mujoco_rs/wrappers/mj_data.rs

mujoco_rs/wrappers/
mj_data.rs
