From 2172383c8eaed45060b4950eb63dd1dabfff1d2a Mon Sep 17 00:00:00 2001 From: Reece Humphreys Date: Sun, 10 May 2026 20:26:56 -0700 Subject: [PATCH 1/2] [#1392] Split qpos/qvel into translation/attitude StateData --- .../_GeneralModuleFiles/MJJoint.cpp | 55 +-- .../MJQPosAttitudeStateData.cpp | 68 ++++ .../MJQPosAttitudeStateData.h | 84 +++++ .../_GeneralModuleFiles/MJScene.cpp | 332 ++++++++++++++++-- .../_GeneralModuleFiles/MJScene.h | 114 +++++- .../_GeneralModuleFiles/MJSpec.cpp | 13 +- 6 files changed, 596 insertions(+), 70 deletions(-) create mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp create mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp index 58c47907d14..b288c5e6f60 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp @@ -102,31 +102,33 @@ void MJScalarJoint::writeJointStateMessage(uint64_t CurrentSimNanos) { checkInitialized(); + auto& scene = body.getSpec().getScene(); + ScalarJointStateMsgPayload stateOutMsgPayload; - auto state = body.getSpec().getScene().getQposState()->getState(); - stateOutMsgPayload.state = state(this->qposAdr.value()); - this->stateOutMsg.write(&stateOutMsgPayload, body.getSpec().getScene().moduleID, 0); + stateOutMsgPayload.state = scene.qposGet(this->qposAdr.value()); + this->stateOutMsg.write(&stateOutMsgPayload, scene.moduleID, 0); ScalarJointStateMsgPayload stateDotOutMsgPayload; - auto stateDot = body.getSpec().getScene().getQvelState()->getState(); - stateDotOutMsgPayload.state = stateDot(this->qvelAdr.value()); - this->stateDotOutMsg.write(&stateDotOutMsgPayload, body.getSpec().getScene().moduleID, 0); + stateDotOutMsgPayload.state = scene.qvelGet(this->qvelAdr.value()); + this->stateDotOutMsg.write(&stateDotOutMsgPayload, scene.moduleID, 0); } void MJScalarJoint::setPosition(double value) { checkInitialized(); - body.getSpec().getScene().getQposState()->state(this->qposAdr.value()) = value; - body.getSpec().getScene().markKinematicsAsStale(); + auto& scene = body.getSpec().getScene(); + scene.qposRef(this->qposAdr.value()) = value; + scene.markKinematicsAsStale(); } void MJScalarJoint::setVelocity(double value) { checkInitialized(); - body.getSpec().getScene().getQvelState()->state(this->qvelAdr.value()) = value; - body.getSpec().getScene().markKinematicsAsStale(); + auto& scene = body.getSpec().getScene(); + scene.qvelRef(this->qvelAdr.value()) = value; + scene.markKinematicsAsStale(); } MJSingleJointEquality @@ -139,45 +141,54 @@ void MJFreeJoint::setPosition(const Eigen::Vector3d& position) { checkInitialized(); - auto& qposState = body.getSpec().getScene().getQposState()->state; + auto& scene = body.getSpec().getScene(); auto i = this->qposAdr.value(); - qposState.middleRows(i, 3) = position; + scene.qposRef(i + 0) = position[0]; + scene.qposRef(i + 1) = position[1]; + scene.qposRef(i + 2) = position[2]; - body.getSpec().getScene().markKinematicsAsStale(); + scene.markKinematicsAsStale(); } void MJFreeJoint::setVelocity(const Eigen::Vector3d& velocity) { checkInitialized(); - auto& qvelState = body.getSpec().getScene().getQvelState()->state; + auto& scene = body.getSpec().getScene(); auto i = this->qvelAdr.value(); - qvelState.middleRows(i, 3) = velocity; + scene.qvelRef(i + 0) = velocity[0]; + scene.qvelRef(i + 1) = velocity[1]; + scene.qvelRef(i + 2) = velocity[2]; - body.getSpec().getScene().markKinematicsAsStale(); + scene.markKinematicsAsStale(); } void MJFreeJoint::setAttitude(const Eigen::MRPd& attitude) { checkInitialized(); - auto& qposState = body.getSpec().getScene().getQposState()->state; + auto& scene = body.getSpec().getScene(); auto i = this->qposAdr.value(); auto mat = attitude.toRotationMatrix(); auto quat = Eigen::Quaterniond(mat); - qposState.middleRows(i + 3, 4) = Eigen::Vector4d{quat.w(), quat.x(), quat.y(), quat.z()}; - body.getSpec().getScene().markKinematicsAsStale(); + scene.qposRef(i + 3) = quat.w(); + scene.qposRef(i + 4) = quat.x(); + scene.qposRef(i + 5) = quat.y(); + scene.qposRef(i + 6) = quat.z(); + scene.markKinematicsAsStale(); } void MJFreeJoint::setAttitudeRate(const Eigen::Vector3d& attitudeRate) { checkInitialized(); - auto& qvelState = body.getSpec().getScene().getQvelState()->state; + auto& scene = body.getSpec().getScene(); auto i = this->qvelAdr.value(); - qvelState.middleRows(i + 3, 3) = attitudeRate; + scene.qvelRef(i + 3) = attitudeRate[0]; + scene.qvelRef(i + 4) = attitudeRate[1]; + scene.qvelRef(i + 5) = attitudeRate[2]; - body.getSpec().getScene().markKinematicsAsStale(); + scene.markKinematicsAsStale(); } diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp new file mode 100644 index 00000000000..b95308750ca --- /dev/null +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp @@ -0,0 +1,68 @@ +/* + ISC License + + Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#include "MJQPosAttitudeStateData.h" + +#include + +#include + +std::unique_ptr MJQPosAttitudeStateData::clone() const +{ + return std::unique_ptr(); +} + +void MJQPosAttitudeStateData::configure(int numQuaternions, int numLinearDofs) +{ + this->numQuaternions = numQuaternions; + this->numLinearDofs = numLinearDofs; + + int stateSize = 4 * numQuaternions + numLinearDofs; + int derivSize = 3 * numQuaternions + numLinearDofs; + + this->state.resize(stateSize, 1); + this->stateDeriv.resize(derivSize, 1); +} + +void MJQPosAttitudeStateData::propagateState(double dt, std::vector pseudoStep) +{ + // Quaternion blocks: state has 4 elements per quaternion, stateDeriv + // has 3 elements per angular velocity. + for (int i = 0; i < numQuaternions; ++i) { + mju_quatIntegrate(this->state.data() + 4 * i, + this->stateDeriv.data() + 3 * i, + dt); + } + + // Linear DOFs follow the quaternion block in both vectors. + int linearStart_state = 4 * numQuaternions; + int linearStart_deriv = 3 * numQuaternions; + for (int i = 0; i < numLinearDofs; ++i) { + this->state(linearStart_state + i) += + this->stateDeriv(linearStart_deriv + i) * dt; + } + + if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) { + auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " + + "the integrator tried to propagate it without pseudoSteps. Are you sure " + + "you are using a stochastic integrator?"; + bskLogger.bskLog(BSK_ERROR, errorMsg.c_str()); + throw std::invalid_argument(errorMsg); + } +} diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h new file mode 100644 index 00000000000..d9150c8e471 --- /dev/null +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h @@ -0,0 +1,84 @@ +/* + ISC License + + Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#ifndef MJQPOS_ATTITUDE_STATE_DATA_H +#define MJQPOS_ATTITUDE_STATE_DATA_H + +#include "simulation/dynamics/_GeneralModuleFiles/stateData.h" + +/** + * @brief Position state data for the attitude/hinges portion of an `MJScene`. + * + * Holds the rotational and scalar-joint position DOFs of an MJScene's qpos + * vector, segregated from the free-joint translational DOFs (which live in a + * separate plain `StateData`). The split lets the adaptive integrator scale + * tolerances independently for translation (orbital scales) and attitude + * (radian-scale): without it, the orbital state norm would dominate the + * tolerance computation and either force absurdly small steps everywhere or + * (with looser relTol) allow steps too large for the stiff panel-joint ODE. + * + * Layout convention (assembled by the owning `MJScene`): + * + * * `state` [4 * numQuaternions + numLinearDofs] + * - First `4 * numQuaternions` entries: 4-element unit quaternions, one + * per free-joint rotation or ball joint. + * - Remaining `numLinearDofs` entries: scalar joint positions (hinges + * and slides), one entry each. + * * `stateDeriv` [3 * numQuaternions + numLinearDofs] + * - First `3 * numQuaternions` entries: 3-element angular velocities + * (one per quaternion block, in the same order). + * - Remaining `numLinearDofs` entries: scalar joint velocities. + * + * `propagateState` integrates the quaternion blocks with `mju_quatIntegrate` + * (proper SO(3) update with renormalization) and the linear DOFs with the + * default Euler step `state += stateDeriv * dt`. + */ +class MJQPosAttitudeStateData : public StateData +{ +public: + /** + * @brief Constructs an `MJQPosAttitudeStateData` with the given name and + * initial state matrix (typically empty until `configure` runs). + */ + MJQPosAttitudeStateData(std::string inName, const Eigen::MatrixXd& newState) + : StateData(std::move(inName), newState) {} + + /** Clone constructor (returns null to match the existing MuJoCo state convention). */ + std::unique_ptr clone() const override; + + /** + * @brief Sizes the `state` and `stateDeriv` vectors and records the layout. + * + * @param numQuaternions Number of 4-element quaternion blocks (free + ball joints). + * @param numLinearDofs Number of 1-element scalar-joint DOFs (hinges + slides). + */ + void configure(int numQuaternions, int numLinearDofs); + + /** + * @brief Propagates the state by `dt`. Quaternions are updated with + * `mju_quatIntegrate` and linear DOFs with `state += stateDeriv * dt`. + */ + void propagateState(double dt, std::vector pseudoStep = {}) override; + +protected: + int numQuaternions = 0; + int numLinearDofs = 0; +}; + +#endif diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp index 4c95fcfcbb9..4309d93003a 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp @@ -95,12 +95,20 @@ void MJScene::Reset(uint64_t CurrentSimNanos) void MJScene::initializeDynamics() { - // We need to use a special StateData type for qpos - // since it is integrated using a mujoco function - this->qposState = this->dynManager.registerState(1, 1, "mujocoQpos"); - - this->qvelState = this->dynManager.registerState(1, 1, "mujocoQvel"); - this->actState = this->dynManager.registerState(1, 1, "mujocoAct"); + // Split MuJoCo's qpos/qvel into four states so the adaptive integrator + // can scale tolerances independently for free-joint translation + // (orbital metres / metres-per-second) and free-joint rotation + + // hinges/slides (radians / radians-per-second). A single combined + // state mixes those scales: the orbital component dominates the state + // norm and forces the relative-tolerance band so wide that errors in + // the stiff hinge ODE pass through unflagged, which is what produced + // the spurious panel spin we hit earlier. + this->qposTranslationState = this->dynManager.registerState(1, 1, "mujocoQposTranslation"); + this->qposAttitudeState = this->dynManager.registerState( + 1, 1, "mujocoQposAttitude"); + this->qvelTranslationState = this->dynManager.registerState(1, 1, "mujocoQvelTranslation"); + this->qvelAttitudeState = this->dynManager.registerState(1, 1, "mujocoQvelAttitude"); + this->actState = this->dynManager.registerState(1, 1, "mujocoAct"); for (auto&& body : this->spec.getBodies()) { body.registerStates(DynParamRegisterer( @@ -117,6 +125,11 @@ void MJScene::initializeDynamics() this->spec.configure(); } + // Now that the MuJoCo model is compiled, joint addresses are known + // and we can build the bidirectional maps between full mjData indices + // and split-state indices. + this->rebuildStateIndexMaps(); + // Register the states of the models in the dynamics task std::unordered_set alreadyRegisteredModels; auto registerStatesOnSysModel = [this, &alreadyRegisteredModels](SysModel* sysModelPtr) @@ -224,12 +237,57 @@ void MJScene::equationsOfMotion(double t, double timeStep) "s in MJScene with ID: " + std::to_string(moduleID)); } - // The derivative of the position is simply the state of the velocity - this->qposState->setDerivative(this->qvelState->getState()); + // Replace the floating-frame translational acceleration for each free joint. + // In the co-moving frame, MuJoCo solves the dynamics with the body at the + // origin, so the qacc translation it produces is just floating-point noise + // from the co-moving cancellation — keeping it would drive the RKF45 + // step-size controller off the attitude PID timescale. We therefore + // overwrite it with the true inertial translational acceleration: the sum + // of gravity accelerations from all of the body's gravity sources evaluated + // at the body's saved (un-zeroed) inertial position, or zero if the body + // has no gravity sources (the original zero-gravity behaviour). + for (auto& sf : savedFrames) { + int i = sf.fj->getQvelAdr(); + Eigen::Vector3d gravAccel = sf.fj->getBody().computeGravityAt(sf.r); + mujocoData->qacc[i] = gravAccel[0]; + mujocoData->qacc[i + 1] = gravAccel[1]; + mujocoData->qacc[i + 2] = gravAccel[2]; + } + + // Push computed velocities and accelerations into the split-state + // derivatives. For qpos states the time derivative is the corresponding + // qvel state (which the attitude state interprets as angular velocity + // for quaternion blocks); for qvel states it is the qacc just computed. + { + Eigen::VectorXd dQposTran(this->qposTranslationIdx.size()); + for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { + // Translation qpos derivative == translation qvel. + dQposTran[k] = mujocoData->qvel[this->qvelTranslationIdx[k]]; + } + this->qposTranslationState->setDerivative(dQposTran); + + Eigen::VectorXd dQposAtt(this->qvelAttitudeIdx.size()); + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + // Attitude qpos stateDeriv is sized as qvel-space (3 angular + // velocities per quaternion + 1 per linear DOF); MJQPosAttitudeStateData + // consumes it via mju_quatIntegrate (for quaternion blocks) and + // Euler step (for linear DOFs). + dQposAtt[k] = mujocoData->qvel[this->qvelAttitudeIdx[k]]; + } + this->qposAttitudeState->setDerivative(dQposAtt); + + Eigen::VectorXd dQvelTran(this->qvelTranslationIdx.size()); + for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { + dQvelTran[k] = qacc[this->qvelTranslationIdx[k]]; + } + this->qvelTranslationState->setDerivative(dQvelTran); - // Copy the computed accelerations into the derivative of the velocity - auto qvelDeriv = this->qvelState->stateDeriv.data(); - std::copy_n(qacc, this->spec.getMujocoModel()->nv, qvelDeriv); + Eigen::VectorXd dQvelAtt(this->qvelAttitudeIdx.size()); + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + dQvelAtt[k] = qacc[this->qvelAttitudeIdx[k]]; + } + this->qvelAttitudeState->setDerivative(dQvelAtt); + } // Also copy the derivative of the actuator states, if we have them if (this->spec.getMujocoModel()->na > 0) { @@ -300,31 +358,88 @@ void MJScene::saveToFile(std::string filename) } } -MJQPosStateData* MJScene::getQposState() +StateData* +MJScene::getQposTranslationState() +{ + if (!this->qposTranslationState) { + logAndThrow("Tried to get qpos translation state before initialization."); + } + return this->qposTranslationState; +} + +MJQPosAttitudeStateData* +MJScene::getQposAttitudeState() +{ + if (!this->qposAttitudeState) { + logAndThrow("Tried to get qpos attitude state before initialization."); + } + return this->qposAttitudeState; +} + +StateData* +MJScene::getQvelTranslationState() { - if (!this->qposState) { - logAndThrow("Tried to get qpos state before initialization."); + if (!this->qvelTranslationState) { + logAndThrow("Tried to get qvel translation state before initialization."); } - return this->qposState; + return this->qvelTranslationState; } -StateData* MJScene::getQvelState() +StateData* +MJScene::getQvelAttitudeState() { - if (!this->qvelState) { - logAndThrow("Tried to get qvel state before initialization."); + if (!this->qvelAttitudeState) { + logAndThrow("Tried to get qvel attitude state before initialization."); } - return this->qvelState; + return this->qvelAttitudeState; } StateData* MJScene::getActState() { if (!this->actState) { - logAndThrow("Tried to get qpos state before initialization."); + logAndThrow("Tried to get act state before initialization."); } return this->actState; } -void MJScene::printMujocoModelDebugInfo(const std::string& path) +double& +MJScene::qposRef(int fullQposIdx) +{ + int splitIdx = this->fullQposToSplitIdx.at(fullQposIdx); + return this->fullQposIsTranslation.at(fullQposIdx) + ? this->qposTranslationState->state(splitIdx) + : this->qposAttitudeState->state(splitIdx); +} + +double +MJScene::qposGet(int fullQposIdx) const +{ + int splitIdx = this->fullQposToSplitIdx.at(fullQposIdx); + return this->fullQposIsTranslation.at(fullQposIdx) + ? this->qposTranslationState->state(splitIdx) + : this->qposAttitudeState->state(splitIdx); +} + +double& +MJScene::qvelRef(int fullQvelIdx) +{ + int splitIdx = this->fullQvelToSplitIdx.at(fullQvelIdx); + return this->fullQvelIsTranslation.at(fullQvelIdx) + ? this->qvelTranslationState->state(splitIdx) + : this->qvelAttitudeState->state(splitIdx); +} + +double +MJScene::qvelGet(int fullQvelIdx) const +{ + int splitIdx = this->fullQvelToSplitIdx.at(fullQvelIdx); + return this->fullQvelIsTranslation.at(fullQvelIdx) + ? this->qvelTranslationState->state(splitIdx) + : this->qvelAttitudeState->state(splitIdx); +} + +void +MJScene::printMujocoModelDebugInfo(const std::string& path) { mj_printModel(this->getMujocoModel(), path.c_str()); } @@ -447,11 +562,25 @@ void MJScene::updateMujocoArraysFromStates() auto mujocoModel = this->getMujocoModel(); auto mujocoData = this->getMujocoData(); - // Copy the joint positions and velocity stored in the states - // into the mujoco arrays - std::copy_n(this->qposState->state.data(), mujocoModel->nq, mujocoData->qpos); + // Reassemble mjData::qpos from the two split position states. + for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { + mujocoData->qpos[this->qposTranslationIdx[k]] = + this->qposTranslationState->state(k); + } + for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { + mujocoData->qpos[this->qposAttitudeIdx[k]] = + this->qposAttitudeState->state(k); + } - std::copy_n(this->qvelState->state.data(), mujocoModel->nv, mujocoData->qvel); + // Reassemble mjData::qvel from the two split velocity states. + for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { + mujocoData->qvel[this->qvelTranslationIdx[k]] = + this->qvelTranslationState->state(k); + } + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + mujocoData->qvel[this->qvelAttitudeIdx[k]] = + this->qvelAttitudeState->state(k); + } if (mujocoModel->na > 0) { std::copy_n(this->actState->state.data(), mujocoModel->na, mujocoData->act); @@ -460,11 +589,156 @@ void MJScene::updateMujocoArraysFromStates() markKinematicsAsStale(); } -void MJScene::writeOutputStateMessages(uint64_t CurrentSimNanos) +Eigen::VectorXd +MJScene::assembleFullQpos() { - MJSceneStateMsgPayload stateOutMsgPayload{this->qposState->getState(), - this->qvelState->getState(), - this->actState->getState()}; + auto m = this->spec.getMujocoModel(); + Eigen::VectorXd qpos = Eigen::VectorXd::Zero(m->nq); + for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { + qpos(this->qposTranslationIdx[k]) = this->qposTranslationState->state(k); + } + for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { + qpos(this->qposAttitudeIdx[k]) = this->qposAttitudeState->state(k); + } + return qpos; +} + +Eigen::VectorXd +MJScene::assembleFullQvel() +{ + auto m = this->spec.getMujocoModel(); + Eigen::VectorXd qvel = Eigen::VectorXd::Zero(m->nv); + for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { + qvel(this->qvelTranslationIdx[k]) = this->qvelTranslationState->state(k); + } + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + qvel(this->qvelAttitudeIdx[k]) = this->qvelAttitudeState->state(k); + } + return qvel; +} + +void +MJScene::writeOutputStateMessages(uint64_t CurrentSimNanos) +{ + MJSceneStateMsgPayload stateOutMsgPayload{ this->assembleFullQpos(), + this->assembleFullQvel(), + this->actState->getState() }; stateOutMsg.write(&stateOutMsgPayload, this->moduleID, CurrentSimNanos); } + +void +MJScene::rebuildStateIndexMaps() +{ + auto m = this->spec.getMujocoModel(); + + // Categorise joints by type using MuJoCo's compiled tables (instead of + // walking the MJBody hierarchy) so we get joint addresses in the same + // order MuJoCo lays them out in qpos/qvel. + std::vector freeJointIds; + std::vector ballJointIds; + std::vector scalarJointIds; + for (int j = 0; j < m->njnt; ++j) { + switch (m->jnt_type[j]) { + case mjJNT_FREE: freeJointIds.push_back(j); break; + case mjJNT_BALL: ballJointIds.push_back(j); break; + case mjJNT_HINGE: + case mjJNT_SLIDE: scalarJointIds.push_back(j); break; + } + } + + this->qposTranslationIdx.clear(); + this->qvelTranslationIdx.clear(); + this->qposAttitudeIdx.clear(); + this->qvelAttitudeIdx.clear(); + + // Translation state: free-joint translational DOFs only. + for (int jid : freeJointIds) { + int qposAdr = m->jnt_qposadr[jid]; + int qvelAdr = m->jnt_dofadr[jid]; + for (int k = 0; k < 3; ++k) { + this->qposTranslationIdx.push_back(qposAdr + k); + this->qvelTranslationIdx.push_back(qvelAdr + k); + } + } + + // Attitude state — quaternion blocks first (free-joint rotation, then + // ball joints), then scalar linear DOFs. This ordering matches + // MJQPosAttitudeStateData::propagateState, which expects all + // quaternions packed together at the front. + for (int jid : freeJointIds) { + int qposAdr = m->jnt_qposadr[jid]; + int qvelAdr = m->jnt_dofadr[jid]; + for (int k = 0; k < 4; ++k) this->qposAttitudeIdx.push_back(qposAdr + 3 + k); + for (int k = 0; k < 3; ++k) this->qvelAttitudeIdx.push_back(qvelAdr + 3 + k); + } + for (int jid : ballJointIds) { + int qposAdr = m->jnt_qposadr[jid]; + int qvelAdr = m->jnt_dofadr[jid]; + for (int k = 0; k < 4; ++k) this->qposAttitudeIdx.push_back(qposAdr + k); + for (int k = 0; k < 3; ++k) this->qvelAttitudeIdx.push_back(qvelAdr + k); + } + for (int jid : scalarJointIds) { + this->qposAttitudeIdx.push_back(m->jnt_qposadr[jid]); + this->qvelAttitudeIdx.push_back(m->jnt_dofadr[jid]); + } + + int numQuaternions = static_cast(freeJointIds.size() + ballJointIds.size()); + int numLinearDofs = static_cast(scalarJointIds.size()); + + // Size the split states. qposAttitudeState has separate state/deriv + // sizes (4 vs 3 per quaternion); the others use the default StateData + // layout where state and stateDeriv are the same size. + this->qposAttitudeState->configure(numQuaternions, numLinearDofs); + + auto resizeMatching = [](StateData* s, int n) { + s->state.resize(n, 1); + s->stateDeriv.resize(n, 1); + }; + resizeMatching(this->qposTranslationState, static_cast(this->qposTranslationIdx.size())); + resizeMatching(this->qvelTranslationState, static_cast(this->qvelTranslationIdx.size())); + resizeMatching(this->qvelAttitudeState, static_cast(this->qvelAttitudeIdx.size())); + + // Reverse maps for joint-side accessors (qposRef/qvelRef). + this->fullQposToSplitIdx.assign(m->nq, -1); + this->fullQposIsTranslation.assign(m->nq, false); + for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { + int fullIdx = this->qposTranslationIdx[k]; + this->fullQposToSplitIdx[fullIdx] = static_cast(k); + this->fullQposIsTranslation[fullIdx] = true; + } + for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { + int fullIdx = this->qposAttitudeIdx[k]; + this->fullQposToSplitIdx[fullIdx] = static_cast(k); + this->fullQposIsTranslation[fullIdx] = false; + } + this->fullQvelToSplitIdx.assign(m->nv, -1); + this->fullQvelIsTranslation.assign(m->nv, false); + for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { + int fullIdx = this->qvelTranslationIdx[k]; + this->fullQvelToSplitIdx[fullIdx] = static_cast(k); + this->fullQvelIsTranslation[fullIdx] = true; + } + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + int fullIdx = this->qvelAttitudeIdx[k]; + this->fullQvelToSplitIdx[fullIdx] = static_cast(k); + this->fullQvelIsTranslation[fullIdx] = false; + } + + // Seed the split states from the freshly compiled mjData so the user + // sees the qpos/qvel values declared in the XML before they make any + // setPosition/setVelocity calls. + auto d = this->spec.getMujocoData(); + for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { + this->qposTranslationState->state(k) = d->qpos[this->qposTranslationIdx[k]]; + } + for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { + this->qposAttitudeState->state(k) = d->qpos[this->qposAttitudeIdx[k]]; + } + for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { + this->qvelTranslationState->state(k) = d->qvel[this->qvelTranslationIdx[k]]; + } + for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { + this->qvelAttitudeState->state(k) = d->qvel[this->qvelAttitudeIdx[k]]; + } +} diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h index f9035338ad7..516d525806c 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -28,9 +28,11 @@ #include "MJUtils.h" #include "architecture/_GeneralModuleFiles/sys_model_task.h" #include "architecture/msgPayloadDefCpp/MJSceneStateMsgPayload.h" -#include "MJQPosStateData.h" +#include "MJQPosAttitudeStateData.h" #include "simulation/dynamics/_GeneralModuleFiles/dynamicObject.h" +#include + /** * @brief Represents a dynamic object that solves multi-body dynamics through MuJoCo. * @@ -567,20 +569,53 @@ class MJScene : public DynamicObject mjData* getMujocoData() { return this->spec.getMujocoData(); } /** - * @brief Retrieves the position state data. + * @brief Retrieves the free-joint translational position state data. + * + * Holds the 3-DOF translational positions of every free joint, packed + * contiguously. Lives in its own state (separate from + * `getQposAttitudeState`) so the adaptive integrator can scale tolerances + * for orbital position separately from radian-scale attitude/hinge DOFs. * - * @return Pointer to `MJQPosStateData`. + * @return Pointer to the translational `StateData`. * @throw std::runtime_error If the state has not been initialized. */ - MJQPosStateData* getQposState(); + StateData* getQposTranslationState(); /** - * @brief Retrieves the velocity state data. + * @brief Retrieves the attitude/hinges position state data. * - * @return Pointer to `StateData`. + * Holds free-joint quaternions, ball-joint quaternions, and all scalar + * (hinge/slide) joint positions. Quaternion blocks are integrated with + * `mju_quatIntegrate`; linear DOFs use the default Euler step. + * + * @return Pointer to the attitude `MJQPosAttitudeStateData`. + * @throw std::runtime_error If the state has not been initialized. + */ + MJQPosAttitudeStateData* getQposAttitudeState(); + + /** + * @brief Retrieves the free-joint translational velocity state data. + * + * Holds the 3-DOF translational velocities of every free joint, packed + * contiguously, separate from attitude/hinge velocities for tolerance + * scaling. + * + * @return Pointer to the translational velocity `StateData`. + * @throw std::runtime_error If the state has not been initialized. + */ + StateData* getQvelTranslationState(); + + /** + * @brief Retrieves the attitude/hinges velocity state data. + * + * Holds free-joint angular velocities, ball-joint angular velocities, + * and all scalar joint velocities. It is the integrator-side derivative + * counterpart of `getQposAttitudeState`. + * + * @return Pointer to the attitude velocity `StateData`. * @throw std::runtime_error If the state has not been initialized. */ - StateData* getQvelState(); + StateData* getQvelAttitudeState(); /** * @brief Retrieves the actuation state data. @@ -590,6 +625,26 @@ class MJScene : public DynamicObject */ StateData* getActState(); + /** + * @brief Returns a reference to the qpos element at the given full + * (mjData::qpos) index, dispatching to the translation or attitude split + * state. Used by joint setters to mutate state by full mujoco-data + * address without knowing which split state owns the element. + * + * @throw std::out_of_range if the index has not been mapped (caller + * passed an invalid index or `initializeDynamics` did not run). + */ + double& qposRef(int fullQposIdx); + + /** Const-read variant of `qposRef`. */ + double qposGet(int fullQposIdx) const; + + /** Same as `qposRef` but for the qvel/qacc indexing space. */ + double& qvelRef(int fullQvelIdx); + + /** Const-read variant of `qvelRef`. */ + double qvelGet(int fullQvelIdx) const; + /** * @brief Prints MuJoCo model debug information to a file. * @@ -624,6 +679,17 @@ class MJScene : public DynamicObject Message stateOutMsg; ///< Message with all the the scene's position, velocity, and actuators states. + /** + * @brief Returns the total qpos vector reassembled from the split states. + * + * This matches the layout of `mjData::qpos` (length `nq`). Useful for + * callers (e.g. the state output message) that want the unified view. + */ + Eigen::VectorXd assembleFullQpos(); + + /** Companion to `assembleFullQpos`, for the qvel space (length `nv`). */ + Eigen::VectorXd assembleFullQvel(); + protected: /** * @brief Updates MuJoCo structs from the Basilisk `StateData` objects. @@ -637,6 +703,14 @@ class MJScene : public DynamicObject */ void writeOutputStateMessages(uint64_t CurrentSimNanos); + /** + * @brief Builds the index maps that translate full mjData qpos/qvel + * indices to/from indices in the split state vectors. Run once per + * `initializeDynamics`, after the MuJoCo spec has been compiled and + * joint addresses are known. + */ + void rebuildStateIndexMaps(); + protected: MJSpec spec; ///< `MJSpec` (MuJoCo model specification wrapper) associated with this scene. bool mjModelConstStale = false; ///< Flag indicating stale model constants. @@ -646,9 +720,29 @@ class MJScene : public DynamicObject SysModelTask dynamicsDiffusionTask; ///< Task managing models involved in the diffusion stochastic dynamics of this scene. std::vector> ownedSysModel; ///< System models that should be cleared on this scene destruction. - MJQPosStateData* qposState; ///< Position state data. - StateData* qvelState; ///< Velocity state data. - StateData* actState; ///< Actuator state data. + StateData* qposTranslationState = nullptr; ///< Free-joint translational positions (3 per free joint). + MJQPosAttitudeStateData* qposAttitudeState = nullptr; ///< Quaternions (free + ball joints) and scalar joint positions. + StateData* qvelTranslationState = nullptr; ///< Free-joint translational velocities (3 per free joint). + StateData* qvelAttitudeState = nullptr; ///< Angular velocities (free + ball) and scalar joint velocities. + StateData* actState = nullptr; ///< Actuator state data. + + /** mjData::qpos index for each element of qposTranslationState. */ + std::vector qposTranslationIdx; + /** mjData::qvel index for each element of qvelTranslationState. */ + std::vector qvelTranslationIdx; + /** mjData::qpos index for each element of qposAttitudeState. */ + std::vector qposAttitudeIdx; + /** mjData::qvel index for each element of qvelAttitudeState. */ + std::vector qvelAttitudeIdx; + + /** Reverse map: full qpos index => index in its split state. */ + std::vector fullQposToSplitIdx; + /** Reverse map: full qpos index -> true if owned by translation state. */ + std::vector fullQposIsTranslation; + /** Reverse map for the qvel space. */ + std::vector fullQvelToSplitIdx; + /** Reverse map for the qvel space (true if translation). */ + std::vector fullQvelIsTranslation; }; template diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp index 7fa4e9a9669..c624772da98 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp @@ -293,15 +293,10 @@ void MJSpec::configure() equality.configure(this->model.get()); } - // Update the sizes of the states according to the model size - // and copy the state in mjData to the Basilisk state objects - this->scene.getQposState()->configure(this->model.get()); - std::copy_n(this->data->qpos, this->model->nq, this->scene.getQposState()->state.data()); - - this->scene.getQvelState()->state.resize(this->model->nv, 1); - this->scene.getQvelState()->stateDeriv.resize(this->model->nv, 1); - std::copy_n(this->data->qvel, this->model->nv, this->scene.getQvelState()->state.data()); - + // Sizing/seeding the qpos/qvel split states from mjData is done by + // MJScene::rebuildStateIndexMaps, which is invoked by + // MJScene::initializeDynamics after this configure() returns. Only the + // act state can be sized here because it's a flat copy with no split. this->scene.getActState()->state.resize(this->model->na, 1); this->scene.getActState()->stateDeriv.resize(this->model->na, 1); std::copy_n(this->data->act, this->model->na, this->scene.getActState()->state.data()); From 0d3239f1a82e9eeabc1b34ee4827bc67e0b16c18 Mon Sep 17 00:00:00 2001 From: Reece Humphreys Date: Wed, 13 May 2026 23:07:21 -0700 Subject: [PATCH 2/2] [#0000] SQUASH: juans proposal --- .../quaternionStateData.cpp | 97 ++++++ .../_GeneralModuleFiles/quaternionStateData.h | 59 ++++ .../_GeneralModuleFiles/MJBody.cpp | 34 ++ .../_GeneralModuleFiles/MJBody.h | 24 +- .../_GeneralModuleFiles/MJBody.swg | 3 + .../_GeneralModuleFiles/MJJoint.cpp | 194 ++++++++--- .../_GeneralModuleFiles/MJJoint.h | 119 ++++--- .../_GeneralModuleFiles/MJJoint.swg | 18 + .../MJQPosAttitudeStateData.cpp | 68 ---- .../MJQPosAttitudeStateData.h | 84 ----- .../_GeneralModuleFiles/MJQPosStateData.cpp | 51 --- .../_GeneralModuleFiles/MJQPosStateData.h | 86 ----- .../_GeneralModuleFiles/MJScene.cpp | 323 ++---------------- .../_GeneralModuleFiles/MJScene.h | 100 ------ .../_GeneralModuleFiles/MJSpec.cpp | 8 +- 15 files changed, 497 insertions(+), 771 deletions(-) create mode 100644 src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.cpp create mode 100644 src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.h delete mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp delete mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h delete mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp delete mode 100644 src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h diff --git a/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.cpp b/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.cpp new file mode 100644 index 00000000000..4f31b90b2e4 --- /dev/null +++ b/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.cpp @@ -0,0 +1,97 @@ +/* + ISC License + + Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#include "quaternionStateData.h" + +#include +#include + +QuaternionStateData::QuaternionStateData(std::string inName, + const Eigen::MatrixXd& newState) + : StateData(std::move(inName), newState) +{ + // state holds (w,x,y,z); deriv is the body angular velocity (3x1). + this->state.resize(4, 1); + this->state.setZero(); + this->state(0) = 1.0; // identity quaternion + this->stateDeriv.resize(3, 1); + this->stateDeriv.setZero(); +} + +std::unique_ptr QuaternionStateData::clone() const +{ + return std::unique_ptr(); +} + +void QuaternionStateData::propagateState(double dt, std::vector pseudoStep) +{ + // Equivalent to mju_quatIntegrate: q <- q * exp(0.5 * dt * omega), normalize. + double wx = this->stateDeriv(0) * dt; + double wy = this->stateDeriv(1) * dt; + double wz = this->stateDeriv(2) * dt; + double angle = std::sqrt(wx * wx + wy * wy + wz * wz); + + double dq_w, dq_x, dq_y, dq_z; + if (angle < 1e-12) { + // Small-angle limit avoids division by zero; sin(a/2)/a -> 1/2. + dq_w = 1.0; + dq_x = 0.5 * wx; + dq_y = 0.5 * wy; + dq_z = 0.5 * wz; + } else { + double half = 0.5 * angle; + double scale = std::sin(half) / angle; + dq_w = std::cos(half); + dq_x = scale * wx; + dq_y = scale * wy; + dq_z = scale * wz; + } + + double qw = this->state(0); + double qx = this->state(1); + double qy = this->state(2); + double qz = this->state(3); + + // Hamilton product q' = q ⊗ dq (body-frame angular velocity convention). + double nw = qw * dq_w - qx * dq_x - qy * dq_y - qz * dq_z; + double nx = qw * dq_x + qx * dq_w + qy * dq_z - qz * dq_y; + double ny = qw * dq_y - qx * dq_z + qy * dq_w + qz * dq_x; + double nz = qw * dq_z + qx * dq_y - qy * dq_x + qz * dq_w; + + double norm = std::sqrt(nw * nw + nx * nx + ny * ny + nz * nz); + if (norm > 0.0) { + double inv = 1.0 / norm; + nw *= inv; nx *= inv; ny *= inv; nz *= inv; + } else { + nw = 1.0; nx = 0.0; ny = 0.0; nz = 0.0; + } + + this->state(0) = nw; + this->state(1) = nx; + this->state(2) = ny; + this->state(3) = nz; + + if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) { + auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " + + "the integrator tried to propagate it without pseudoSteps. Are you sure " + + "you are using a stochastic integrator?"; + bskLogger.bskLog(BSK_ERROR, "%s", errorMsg.c_str()); + throw std::invalid_argument(errorMsg); + } +} diff --git a/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.h b/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.h new file mode 100644 index 00000000000..87e74b52ab2 --- /dev/null +++ b/src/simulation/dynamics/_GeneralModuleFiles/quaternionStateData.h @@ -0,0 +1,59 @@ +/* + ISC License + + Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder + + Permission to use, copy, modify, and/or distribute this software for any + purpose with or without fee is hereby granted, provided that the above + copyright notice and this permission notice appear in all copies. + + THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES + WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF + MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR + ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES + WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN + ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF + OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. + + */ + +#ifndef QUATERNION_STATE_DATA_H +#define QUATERNION_STATE_DATA_H + +#include +#include +#include + +#include + +#include "stateData.h" + +/** + * @brief A `StateData` whose value is a unit quaternion advanced from a body + * angular velocity. + * + * `state` is a 4x1 unit quaternion in `(w, x, y, z)` order. + * `stateDeriv` is a 3x1 body angular velocity in rad/s. + * + * Sizes intentionally differ — the standard Euler step `state += deriv * dt` + * is not valid on SO(3). `propagateState` instead applies the exact + * exponential-map integrator and renormalizes so `|q| = 1` after each step. + */ +class QuaternionStateData : public StateData +{ +public: + /** Constructs a 4x1 quaternion state with a 3x1 derivative. */ + QuaternionStateData(std::string inName, const Eigen::MatrixXd& newState); + + /** Quaternion states are owned and seeded by their joint; cloning returns null. */ + std::unique_ptr clone() const override; + + /** + * @brief Integrates `state` over `dt` using the body angular velocity in + * `stateDeriv`. Equivalent to MuJoCo's `mju_quatIntegrate`: composes the + * current quaternion with `exp(0.5 * dt * omega)` and renormalizes. + */ + void propagateState(double dt, std::vector pseudoStep = {}) override; +}; + +#endif diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp index 224c9725cb9..4782e839868 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.cpp @@ -252,6 +252,40 @@ void MJBody::writeStateDependentOutputMessages(uint64_t CurrentSimNanos) void MJBody::registerStates(DynParamRegisterer paramManager) { this->massState = paramManager.registerState(1, 1, "mass"); + + // Each joint owns its own qpos/qvel state(s). Registering here puts + // them all under this body's prefix so the manager-wide names stay + // unique even when bodies have joints with the same XML name. + for (auto&& joint : this->scalarJoints) { + joint.registerStates(paramManager); + } + if (this->ballJoint.has_value()) { + this->ballJoint->registerStates(paramManager); + } + if (this->freeJoint.has_value()) { + this->freeJoint->registerStates(paramManager); + } +} + +void MJBody::syncJointStatesToMujoco(mjData* d) const +{ + for (auto&& joint : this->scalarJoints) joint.writeStateToMujoco(d); + if (this->ballJoint.has_value()) this->ballJoint->writeStateToMujoco(d); + if (this->freeJoint.has_value()) this->freeJoint->writeStateToMujoco(d); +} + +void MJBody::seedJointStatesFromMujoco(const mjData* d) +{ + for (auto&& joint : this->scalarJoints) joint.readStateFromMujoco(d); + if (this->ballJoint.has_value()) this->ballJoint->readStateFromMujoco(d); + if (this->freeJoint.has_value()) this->freeJoint->readStateFromMujoco(d); +} + +void MJBody::setJointDerivativesFromMujoco(const mjData* d) +{ + for (auto&& joint : this->scalarJoints) joint.setDerivativesFromMujoco(d); + if (this->ballJoint.has_value()) this->ballJoint->setDerivativesFromMujoco(d); + if (this->freeJoint.has_value()) this->freeJoint->setDerivativesFromMujoco(d); } void MJBody::updateMujocoModelFromMassProps() diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h index 2d12ac1e4c4..1af87217f51 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.h @@ -253,14 +253,28 @@ class MJBody : public MJObject void writeStateDependentOutputMessages(uint64_t CurrentSimNanos); /** - * @brief Registers the body's states with a dynamic parameter registerer. - * - * Currently, only the mass of the body is considered a parameter. - * - * @param paramManager The dynamic parameter registerer. + * @brief Registers the body's mass state and the qpos/qvel states of all + * joints attached to this body. Each joint owns its own `StateData` so + * the adaptive integrator can scale tolerances independently per joint. */ void registerStates(DynParamRegisterer paramManager); + /** Pushes the qpos/qvel values from each owned joint's state into `mjData`. */ + void syncJointStatesToMujoco(mjData* d) const; + + /** + * Seeds each owned joint's state from `mjData::qpos`/`mjData::qvel`. + * Called once after spec compile so the user-visible state matches the + * values declared in the XML before any setPosition/setVelocity calls. + */ + void seedJointStatesFromMujoco(const mjData* d); + + /** + * Sets each owned joint's qpos/qvel state derivatives from the current + * `mjData::qvel`/`mjData::qacc`. + */ + void setJointDerivativesFromMujoco(const mjData* d); + /** * @brief Updates the MuJoCo model from the mass properties of the body. * diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.swg b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.swg index 21d69b918f6..52eafb87abc 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.swg +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJBody.swg @@ -33,6 +33,9 @@ %ignore MJBody::updateMujocoModelFromMassProps; %ignore MJBody::updateMassPropsDerivative; %ignore MJBody::updateConstrainedEqualityJoints; +%ignore MJBody::syncJointStatesToMujoco; +%ignore MJBody::seedJointStatesFromMujoco; +%ignore MJBody::setJointDerivativesFromMujoco; %extend MJBody { MJScene& getScene() { return $self->getSpec().getScene(); } diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp index b288c5e6f60..3be040222ee 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.cpp @@ -61,6 +61,10 @@ void MJJoint::checkInitialized() const } } +// --------------------------------------------------------------------------- +// MJScalarJoint +// --------------------------------------------------------------------------- + MJScalarJoint::MJScalarJoint(mjsJoint* joint, MJBody& body) : MJJoint(joint, body), constrainedEquality( @@ -98,6 +102,35 @@ void MJScalarJoint::updateConstrainedEquality() } } +void MJScalarJoint::registerStates(DynParamRegisterer registerer) +{ + this->qposState = registerer.registerState(1, 1, "joint_" + this->name + "_qpos"); + this->qvelState = registerer.registerState(1, 1, "joint_" + this->name + "_qvel"); +} + +void MJScalarJoint::writeStateToMujoco(mjData* d) const +{ + d->qpos[this->qposAdr.value()] = this->qposState->state(0); + d->qvel[this->qvelAdr.value()] = this->qvelState->state(0); +} + +void MJScalarJoint::readStateFromMujoco(const mjData* d) +{ + this->qposState->state(0) = d->qpos[this->qposAdr.value()]; + this->qvelState->state(0) = d->qvel[this->qvelAdr.value()]; +} + +void MJScalarJoint::setDerivativesFromMujoco(const mjData* d) +{ + Eigen::Matrix dPos; + dPos(0) = d->qvel[this->qvelAdr.value()]; + this->qposState->setDerivative(dPos); + + Eigen::Matrix dVel; + dVel(0) = d->qacc[this->qvelAdr.value()]; + this->qvelState->setDerivative(dVel); +} + void MJScalarJoint::writeJointStateMessage(uint64_t CurrentSimNanos) { checkInitialized(); @@ -105,30 +138,26 @@ void MJScalarJoint::writeJointStateMessage(uint64_t CurrentSimNanos) auto& scene = body.getSpec().getScene(); ScalarJointStateMsgPayload stateOutMsgPayload; - stateOutMsgPayload.state = scene.qposGet(this->qposAdr.value()); + stateOutMsgPayload.state = this->qposState->state(0); this->stateOutMsg.write(&stateOutMsgPayload, scene.moduleID, 0); ScalarJointStateMsgPayload stateDotOutMsgPayload; - stateDotOutMsgPayload.state = scene.qvelGet(this->qvelAdr.value()); + stateDotOutMsgPayload.state = this->qvelState->state(0); this->stateDotOutMsg.write(&stateDotOutMsgPayload, scene.moduleID, 0); } void MJScalarJoint::setPosition(double value) { checkInitialized(); - - auto& scene = body.getSpec().getScene(); - scene.qposRef(this->qposAdr.value()) = value; - scene.markKinematicsAsStale(); + this->qposState->state(0) = value; + this->body.getSpec().getScene().markKinematicsAsStale(); } void MJScalarJoint::setVelocity(double value) { checkInitialized(); - - auto& scene = body.getSpec().getScene(); - scene.qvelRef(this->qvelAdr.value()) = value; - scene.markKinematicsAsStale(); + this->qvelState->state(0) = value; + this->body.getSpec().getScene().markKinematicsAsStale(); } MJSingleJointEquality @@ -137,58 +166,135 @@ MJScalarJoint::getConstrainedEquality() return this->constrainedEquality; } -void MJFreeJoint::setPosition(const Eigen::Vector3d& position) +// --------------------------------------------------------------------------- +// MJBallJoint +// --------------------------------------------------------------------------- + +void MJBallJoint::registerStates(DynParamRegisterer registerer) { - checkInitialized(); + this->qposState = registerer.registerState( + 4, 1, "joint_" + this->name + "_qpos"); + this->qvelState = registerer.registerState(3, 1, "joint_" + this->name + "_qvel"); +} - auto& scene = body.getSpec().getScene(); - auto i = this->qposAdr.value(); - scene.qposRef(i + 0) = position[0]; - scene.qposRef(i + 1) = position[1]; - scene.qposRef(i + 2) = position[2]; +void MJBallJoint::writeStateToMujoco(mjData* d) const +{ + auto qp = this->qposAdr.value(); + auto qv = this->qvelAdr.value(); + for (int k = 0; k < 4; ++k) d->qpos[qp + k] = this->qposState->state(k); + for (int k = 0; k < 3; ++k) d->qvel[qv + k] = this->qvelState->state(k); +} - scene.markKinematicsAsStale(); +void MJBallJoint::readStateFromMujoco(const mjData* d) +{ + auto qp = this->qposAdr.value(); + auto qv = this->qvelAdr.value(); + for (int k = 0; k < 4; ++k) this->qposState->state(k) = d->qpos[qp + k]; + for (int k = 0; k < 3; ++k) this->qvelState->state(k) = d->qvel[qv + k]; } -void MJFreeJoint::setVelocity(const Eigen::Vector3d& velocity) +void MJBallJoint::setDerivativesFromMujoco(const mjData* d) { - checkInitialized(); + auto qv = this->qvelAdr.value(); + // qpos integrator consumes body angular velocity (3 elements). + Eigen::Matrix omega; + for (int k = 0; k < 3; ++k) omega(k) = d->qvel[qv + k]; + this->qposState->setDerivative(omega); + + Eigen::Matrix alpha; + for (int k = 0; k < 3; ++k) alpha(k) = d->qacc[qv + k]; + this->qvelState->setDerivative(alpha); +} - auto& scene = body.getSpec().getScene(); - auto i = this->qvelAdr.value(); - scene.qvelRef(i + 0) = velocity[0]; - scene.qvelRef(i + 1) = velocity[1]; - scene.qvelRef(i + 2) = velocity[2]; +// --------------------------------------------------------------------------- +// MJFreeJoint +// --------------------------------------------------------------------------- - scene.markKinematicsAsStale(); +void MJFreeJoint::registerStates(DynParamRegisterer registerer) +{ + this->qposTranslationState = registerer.registerState( + 3, 1, "joint_" + this->name + "_qposTranslation"); + this->qposAttitudeState = registerer.registerState( + 4, 1, "joint_" + this->name + "_qposAttitude"); + this->qvelTranslationState = registerer.registerState( + 3, 1, "joint_" + this->name + "_qvelTranslation"); + this->qvelAttitudeState = registerer.registerState( + 3, 1, "joint_" + this->name + "_qvelAttitude"); } -void MJFreeJoint::setAttitude(const Eigen::MRPd& attitude) +void MJFreeJoint::writeStateToMujoco(mjData* d) const { - checkInitialized(); + auto qp = this->qposAdr.value(); + auto qv = this->qvelAdr.value(); + for (int k = 0; k < 3; ++k) d->qpos[qp + k] = this->qposTranslationState->state(k); + for (int k = 0; k < 4; ++k) d->qpos[qp + 3 + k] = this->qposAttitudeState->state(k); + for (int k = 0; k < 3; ++k) d->qvel[qv + k] = this->qvelTranslationState->state(k); + for (int k = 0; k < 3; ++k) d->qvel[qv + 3 + k] = this->qvelAttitudeState->state(k); +} - auto& scene = body.getSpec().getScene(); - auto i = this->qposAdr.value(); +void MJFreeJoint::readStateFromMujoco(const mjData* d) +{ + auto qp = this->qposAdr.value(); + auto qv = this->qvelAdr.value(); + for (int k = 0; k < 3; ++k) this->qposTranslationState->state(k) = d->qpos[qp + k]; + for (int k = 0; k < 4; ++k) this->qposAttitudeState->state(k) = d->qpos[qp + 3 + k]; + for (int k = 0; k < 3; ++k) this->qvelTranslationState->state(k) = d->qvel[qv + k]; + for (int k = 0; k < 3; ++k) this->qvelAttitudeState->state(k) = d->qvel[qv + 3 + k]; +} - auto mat = attitude.toRotationMatrix(); - auto quat = Eigen::Quaterniond(mat); +void MJFreeJoint::setDerivativesFromMujoco(const mjData* d) +{ + auto qv = this->qvelAdr.value(); + + // Translation qpos derivative is the inertial translational velocity. + Eigen::Matrix dPosTran; + for (int k = 0; k < 3; ++k) dPosTran(k) = d->qvel[qv + k]; + this->qposTranslationState->setDerivative(dPosTran); + + // Attitude qpos derivative is the body angular velocity (consumed by + // QuaternionStateData::propagateState as omega). + Eigen::Matrix omega; + for (int k = 0; k < 3; ++k) omega(k) = d->qvel[qv + 3 + k]; + this->qposAttitudeState->setDerivative(omega); + + Eigen::Matrix dVelTran; + for (int k = 0; k < 3; ++k) dVelTran(k) = d->qacc[qv + k]; + this->qvelTranslationState->setDerivative(dVelTran); + + Eigen::Matrix dVelAtt; + for (int k = 0; k < 3; ++k) dVelAtt(k) = d->qacc[qv + 3 + k]; + this->qvelAttitudeState->setDerivative(dVelAtt); +} - scene.qposRef(i + 3) = quat.w(); - scene.qposRef(i + 4) = quat.x(); - scene.qposRef(i + 5) = quat.y(); - scene.qposRef(i + 6) = quat.z(); - scene.markKinematicsAsStale(); +void MJFreeJoint::setPosition(const Eigen::Vector3d& position) +{ + checkInitialized(); + for (int k = 0; k < 3; ++k) this->qposTranslationState->state(k) = position[k]; + this->body.getSpec().getScene().markKinematicsAsStale(); } -void MJFreeJoint::setAttitudeRate(const Eigen::Vector3d& attitudeRate) +void MJFreeJoint::setVelocity(const Eigen::Vector3d& velocity) { checkInitialized(); + for (int k = 0; k < 3; ++k) this->qvelTranslationState->state(k) = velocity[k]; + this->body.getSpec().getScene().markKinematicsAsStale(); +} - auto& scene = body.getSpec().getScene(); - auto i = this->qvelAdr.value(); - scene.qvelRef(i + 3) = attitudeRate[0]; - scene.qvelRef(i + 4) = attitudeRate[1]; - scene.qvelRef(i + 5) = attitudeRate[2]; +void MJFreeJoint::setAttitude(const Eigen::MRPd& attitude) +{ + checkInitialized(); + auto mat = attitude.toRotationMatrix(); + auto quat = Eigen::Quaterniond(mat); + this->qposAttitudeState->state(0) = quat.w(); + this->qposAttitudeState->state(1) = quat.x(); + this->qposAttitudeState->state(2) = quat.y(); + this->qposAttitudeState->state(3) = quat.z(); + this->body.getSpec().getScene().markKinematicsAsStale(); +} - scene.markKinematicsAsStale(); +void MJFreeJoint::setAttitudeRate(const Eigen::Vector3d& attitudeRate) +{ + checkInitialized(); + for (int k = 0; k < 3; ++k) this->qvelAttitudeState->state(k) = attitudeRate[k]; + this->body.getSpec().getScene().markKinematicsAsStale(); } diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h index 6e6645f46d8..94e07fb23f6 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.h @@ -26,6 +26,8 @@ #include "architecture/msgPayloadDefC/ScalarJointStateMsgPayload.h" #include "architecture/utilities/avsEigenSupport.h" #include "simulation/dynamics/_GeneralModuleFiles/dynParamManager.h" +#include "simulation/dynamics/_GeneralModuleFiles/quaternionStateData.h" +#include "StatefulSysModel.h" #include "MJEquality.h" #include "MJObject.h" @@ -50,8 +52,10 @@ constexpr std::string_view MJBasilisk::detail::getObjectTypeName() /** * @brief Represents a MuJoCo joint. * - * The `MJJoint` class provides basic functionality to handle joints in MuJoCo. - * This is a base class for more specific joint types. + * Each joint owns its own `StateData` for position and velocity (one or more, + * depending on joint type). This per-joint ownership replaces the bulk + * `mujocoQpos`/`mujocoQvel` arrays so the adaptive integrator computes a + * separate error norm for each joint at its natural scale. */ class MJJoint : public MJObject { @@ -90,6 +94,36 @@ class MJJoint : public MJObject */ void configure(const mjModel* m); + /** + * @brief Registers this joint's qpos/qvel `StateData` on the manager. + * + * Each joint type registers state objects sized for its DOFs (1 for + * scalar, 4-pos/3-vel for ball, 3+4-pos/3+3-vel for free). Names are + * scoped via `registerer`'s prefix so collisions across joints are + * impossible. + */ + virtual void registerStates(DynParamRegisterer registerer) = 0; + + /** + * @brief Copies the joint's owned state values into `mjData::qpos` + * and `mjData::qvel` at this joint's address. + */ + virtual void writeStateToMujoco(mjData* d) const = 0; + + /** + * @brief Reads the joint's slice of `mjData::qpos`/`mjData::qvel` into + * the joint's owned states. Used at initialization to seed states + * from the values declared in the XML. + */ + virtual void readStateFromMujoco(const mjData* d) = 0; + + /** + * @brief Sets the joint's qpos and qvel state derivatives from the + * current `mjData::qvel` (drives qpos integration) and `mjData::qacc` + * (drives qvel integration). + */ + virtual void setDerivativesFromMujoco(const mjData* d) = 0; + protected: /** * @brief Checks if the joint has been properly initialized. @@ -98,6 +132,12 @@ class MJJoint : public MJObject */ void checkInitialized() const; +public: + /** Returns this joint's address into `mjData::qpos`. */ + size_t getQposAdr() const { return qposAdr.value(); } + /** Returns this joint's address into `mjData::qvel`. */ + size_t getQvelAdr() const { return qvelAdr.value(); } + protected: MJBody& body; ///< Reference to the body the joint is attached to. @@ -112,6 +152,8 @@ class MJJoint : public MJObject * degree-of-freedom joints, both linear and angular. The position and velocity * of this joint can be set. * + * Owns a 1x1 `StateData` for position and a 1x1 `StateData` for velocity. + * * If `constrainedStateInMsg` is linked, the value in this message will be read * and applied to an `MJSingleJointEquality` such that the joint is constrained * to hold a specific state. This can be used to enforce joints to follow a @@ -145,62 +187,39 @@ class MJScalarJoint : public MJJoint /** * @brief Sets the position of the joint. - * - * For revolute joints, this is the angle (in radians) of the joint with respect - * to its zero-pose. - * - * For linear joints, this is the displacement (in meters) of the joint with - * respect to its zero-pose. - * - * @param value The desired position value. */ void setPosition(double value); /** * @brief Sets the velocity of the joint. - * - * For revolute joints, this is the angular rate (in radians per second). - * - * For linear joints, this is the displacement rate (in meters per second). - * - * @param value The desired velocity value. */ void setVelocity(double value); /** * @brief Returns the equality constraint object associated with this scalar joint. - * - * The returned MJSingleJointEquality can be used to enforce a specific state - * for the joint within the MuJoCo simulation. This is typically used when the - * joint is constrained to follow a particular position or velocity, as specified - * by an input message or control logic. - * - * @return The MJSingleJointEquality object for this joint. */ MJSingleJointEquality getConstrainedEquality(); /** * @brief Configures the scalar joint within the given MuJoCo model. - * - * @param model Pointer to the MuJoCo model. */ void configure(const mjModel* model); /** * @brief Updates the constrained equality associated with the joint. - * - * If `constrainedStateInMsg` is connected, then activates the equality - * constraint of this joint and sets its 'offset' to the value in the message. */ void updateConstrainedEquality(); /** * @brief Writes the joint state message based on the current simulation state. - * - * @param CurrentSimNanos The current simulation time in nanoseconds. */ void writeJointStateMessage(uint64_t CurrentSimNanos); + void registerStates(DynParamRegisterer registerer) override; + void writeStateToMujoco(mjData* d) const override; + void readStateFromMujoco(const mjData* d) override; + void setDerivativesFromMujoco(const mjData* d) override; + public: Message stateOutMsg; ///< Message to output joint position state. Message stateDotOutMsg; ///< Message to output joint velocity state. @@ -208,6 +227,9 @@ class MJScalarJoint : public MJJoint ReadFunctor constrainedStateInMsg; ///< Functor to read constrained state input. protected: + StateData* qposState = nullptr; ///< 1x1 joint position state. + StateData* qvelState = nullptr; ///< 1x1 joint velocity state. + /** An equality used to enforce a specific state for the joint. */ MJSingleJointEquality constrainedEquality; }; @@ -215,13 +237,23 @@ class MJScalarJoint : public MJJoint /** * @brief Represents a ball joint in MuJoCo. * - * Not fully supported. + * Owns a 4x1 `QuaternionStateData` for the orientation and a 3x1 `StateData` + * for the body angular velocity. Not fully supported elsewhere yet. */ class MJBallJoint : public MJJoint { public: /** Use constructor from MJJoint */ using MJJoint::MJJoint; + + void registerStates(DynParamRegisterer registerer) override; + void writeStateToMujoco(mjData* d) const override; + void readStateFromMujoco(const mjData* d) override; + void setDerivativesFromMujoco(const mjData* d) override; + +protected: + QuaternionStateData* qposState = nullptr; ///< 4x1 quaternion (w,x,y,z). + StateData* qvelState = nullptr; ///< 3x1 body angular velocity. }; /** @@ -230,6 +262,11 @@ class MJBallJoint : public MJJoint * The `MJFreeJoint` class provides additional functionality for joints with three * translational and three rotational degrees of freedom, including setting position, * velocity, attitude, and attitude rate. + * + * Owns four `StateData`s — translation pos/vel (3x1 each) and attitude + * quaternion + angular velocity — split so the adaptive integrator can scale + * tolerances independently for orbital translation (metres) and rotation + * (radians). */ class MJFreeJoint : public MJJoint { @@ -239,32 +276,34 @@ class MJFreeJoint : public MJJoint /** * @brief Sets the position of the free joint. - * - * @param position The desired 3D position vector with respect to the inertial frame. */ void setPosition(const Eigen::Vector3d& position); /** * @brief Sets the translational velocity of the free joint. - * - * @param velocity The desired 3D velocity vector with respect to the inertial frame. */ void setVelocity(const Eigen::Vector3d& velocity); /** * @brief Sets the attitude (orientation) of the free joint. - * - * @param attitude The orientation represented as Modified Rodrigues Parameters (MRP) with respect to the inertial frame. */ void setAttitude(const Eigen::MRPd& attitude); /** * @brief Sets the attitude rate (angular velocity) of the free joint. - * - * @param attitudeRate The desired 3D attitude rate vector with respect to the inertial frame. - * @todo Verify if this matches the expected attitude rate conventions with Basilisk. */ void setAttitudeRate(const Eigen::Vector3d& attitudeRate); + + void registerStates(DynParamRegisterer registerer) override; + void writeStateToMujoco(mjData* d) const override; + void readStateFromMujoco(const mjData* d) override; + void setDerivativesFromMujoco(const mjData* d) override; + +protected: + StateData* qposTranslationState = nullptr; ///< 3x1 inertial position. + QuaternionStateData* qposAttitudeState = nullptr; ///< 4x1 attitude quaternion. + StateData* qvelTranslationState = nullptr; ///< 3x1 inertial velocity. + StateData* qvelAttitudeState = nullptr; ///< 3x1 body angular velocity. }; #endif diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.swg b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.swg index 433bab42461..189f3c04c7d 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.swg +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJJoint.swg @@ -25,11 +25,29 @@ %template_mujoco_object(mjsJoint) %ignore MJJoint::MJJoint; +%ignore MJJoint::registerStates; +%ignore MJJoint::writeStateToMujoco; +%ignore MJJoint::readStateFromMujoco; +%ignore MJJoint::setDerivativesFromMujoco; %ignore MJScalarJoint::MJScalarJoint; %ignore MJScalarJoint::configure; %ignore MJScalarJoint::updateConstrainedEquality; %ignore MJScalarJoint::writeJointStateMessage; +%ignore MJScalarJoint::registerStates; +%ignore MJScalarJoint::writeStateToMujoco; +%ignore MJScalarJoint::readStateFromMujoco; +%ignore MJScalarJoint::setDerivativesFromMujoco; + +%ignore MJBallJoint::registerStates; +%ignore MJBallJoint::writeStateToMujoco; +%ignore MJBallJoint::readStateFromMujoco; +%ignore MJBallJoint::setDerivativesFromMujoco; + +%ignore MJFreeJoint::registerStates; +%ignore MJFreeJoint::writeStateToMujoco; +%ignore MJFreeJoint::readStateFromMujoco; +%ignore MJFreeJoint::setDerivativesFromMujoco; %include "MJJoint.h" diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp deleted file mode 100644 index b95308750ca..00000000000 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.cpp +++ /dev/null @@ -1,68 +0,0 @@ -/* - ISC License - - Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder - - Permission to use, copy, modify, and/or distribute this software for any - purpose with or without fee is hereby granted, provided that the above - copyright notice and this permission notice appear in all copies. - - THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES - WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF - MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR - ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES - WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN - ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF - OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. - - */ - -#include "MJQPosAttitudeStateData.h" - -#include - -#include - -std::unique_ptr MJQPosAttitudeStateData::clone() const -{ - return std::unique_ptr(); -} - -void MJQPosAttitudeStateData::configure(int numQuaternions, int numLinearDofs) -{ - this->numQuaternions = numQuaternions; - this->numLinearDofs = numLinearDofs; - - int stateSize = 4 * numQuaternions + numLinearDofs; - int derivSize = 3 * numQuaternions + numLinearDofs; - - this->state.resize(stateSize, 1); - this->stateDeriv.resize(derivSize, 1); -} - -void MJQPosAttitudeStateData::propagateState(double dt, std::vector pseudoStep) -{ - // Quaternion blocks: state has 4 elements per quaternion, stateDeriv - // has 3 elements per angular velocity. - for (int i = 0; i < numQuaternions; ++i) { - mju_quatIntegrate(this->state.data() + 4 * i, - this->stateDeriv.data() + 3 * i, - dt); - } - - // Linear DOFs follow the quaternion block in both vectors. - int linearStart_state = 4 * numQuaternions; - int linearStart_deriv = 3 * numQuaternions; - for (int i = 0; i < numLinearDofs; ++i) { - this->state(linearStart_state + i) += - this->stateDeriv(linearStart_deriv + i) * dt; - } - - if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) { - auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " - + "the integrator tried to propagate it without pseudoSteps. Are you sure " - + "you are using a stochastic integrator?"; - bskLogger.bskLog(BSK_ERROR, errorMsg.c_str()); - throw std::invalid_argument(errorMsg); - } -} diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h deleted file mode 100644 index d9150c8e471..00000000000 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosAttitudeStateData.h +++ /dev/null @@ -1,84 +0,0 @@ -/* - ISC License - - Copyright (c) 2026, Autonomous Vehicle Systems Lab, University of Colorado at Boulder - - Permission to use, copy, modify, and/or distribute this software for any - purpose with or without fee is hereby granted, provided that the above - copyright notice and this permission notice appear in all copies. - - THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES - WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF - MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR - ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES - WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN - ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF - OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. - - */ - -#ifndef MJQPOS_ATTITUDE_STATE_DATA_H -#define MJQPOS_ATTITUDE_STATE_DATA_H - -#include "simulation/dynamics/_GeneralModuleFiles/stateData.h" - -/** - * @brief Position state data for the attitude/hinges portion of an `MJScene`. - * - * Holds the rotational and scalar-joint position DOFs of an MJScene's qpos - * vector, segregated from the free-joint translational DOFs (which live in a - * separate plain `StateData`). The split lets the adaptive integrator scale - * tolerances independently for translation (orbital scales) and attitude - * (radian-scale): without it, the orbital state norm would dominate the - * tolerance computation and either force absurdly small steps everywhere or - * (with looser relTol) allow steps too large for the stiff panel-joint ODE. - * - * Layout convention (assembled by the owning `MJScene`): - * - * * `state` [4 * numQuaternions + numLinearDofs] - * - First `4 * numQuaternions` entries: 4-element unit quaternions, one - * per free-joint rotation or ball joint. - * - Remaining `numLinearDofs` entries: scalar joint positions (hinges - * and slides), one entry each. - * * `stateDeriv` [3 * numQuaternions + numLinearDofs] - * - First `3 * numQuaternions` entries: 3-element angular velocities - * (one per quaternion block, in the same order). - * - Remaining `numLinearDofs` entries: scalar joint velocities. - * - * `propagateState` integrates the quaternion blocks with `mju_quatIntegrate` - * (proper SO(3) update with renormalization) and the linear DOFs with the - * default Euler step `state += stateDeriv * dt`. - */ -class MJQPosAttitudeStateData : public StateData -{ -public: - /** - * @brief Constructs an `MJQPosAttitudeStateData` with the given name and - * initial state matrix (typically empty until `configure` runs). - */ - MJQPosAttitudeStateData(std::string inName, const Eigen::MatrixXd& newState) - : StateData(std::move(inName), newState) {} - - /** Clone constructor (returns null to match the existing MuJoCo state convention). */ - std::unique_ptr clone() const override; - - /** - * @brief Sizes the `state` and `stateDeriv` vectors and records the layout. - * - * @param numQuaternions Number of 4-element quaternion blocks (free + ball joints). - * @param numLinearDofs Number of 1-element scalar-joint DOFs (hinges + slides). - */ - void configure(int numQuaternions, int numLinearDofs); - - /** - * @brief Propagates the state by `dt`. Quaternions are updated with - * `mju_quatIntegrate` and linear DOFs with `state += stateDeriv * dt`. - */ - void propagateState(double dt, std::vector pseudoStep = {}) override; - -protected: - int numQuaternions = 0; - int numLinearDofs = 0; -}; - -#endif diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp deleted file mode 100644 index c17bcc35916..00000000000 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.cpp +++ /dev/null @@ -1,51 +0,0 @@ -/* - ISC License - - Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder - - Permission to use, copy, modify, and/or distribute this software for any - purpose with or without fee is hereby granted, provided that the above - copyright notice and this permission notice appear in all copies. - - THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES - WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF - MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR - ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES - WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN - ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF - OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. - - */ - -#include "MJQPosStateData.h" - -#include - -std::unique_ptr MJQPosStateData::clone() const { return std::unique_ptr(); } - -void MJQPosStateData::configure(mjModel* mujocoModel) -{ - this->mujocoModel = mujocoModel; - - this->state.resize(mujocoModel->nq, 1); - this->stateDeriv.resize(mujocoModel->nv, 1); -} - -void MJQPosStateData::propagateState(double dt, std::vector pseudoStep) -{ - mj_integratePos(this->mujocoModel, this->state.data(), this->stateDeriv.data(), dt); - - if (getNumNoiseSources() > 0 && pseudoStep.size() == 0) - { - auto errorMsg = "State " + this->getName() + " has stochastic dynamics, but " - + "the integrator tried to propagate it without pseudoSteps. Are you sure " - + "you are using a stochastic integrator?"; - bskLogger.bskLog(BSK_ERROR, "%s", errorMsg.c_str()); - throw std::invalid_argument(errorMsg); - } - - for (size_t i = 0; i < getNumNoiseSources(); i++) - { - mj_integratePos(this->mujocoModel, this->state.data(), this->stateDiffusion.at(i).data(), pseudoStep.at(i)); - } -} diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h deleted file mode 100644 index b25079303cd..00000000000 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJQPosStateData.h +++ /dev/null @@ -1,86 +0,0 @@ -/* - ISC License - - Copyright (c) 2025, Autonomous Vehicle Systems Lab, University of Colorado at Boulder - - Permission to use, copy, modify, and/or distribute this software for any - purpose with or without fee is hereby granted, provided that the above - copyright notice and this permission notice appear in all copies. - - THE SOFTWARE IS PROVIDED "AS IS" AND THE AUTHOR DISCLAIMS ALL WARRANTIES - WITH REGARD TO THIS SOFTWARE INCLUDING ALL IMPLIED WARRANTIES OF - MERCHANTABILITY AND FITNESS. IN NO EVENT SHALL THE AUTHOR BE LIABLE FOR - ANY SPECIAL, DIRECT, INDIRECT, OR CONSEQUENTIAL DAMAGES OR ANY DAMAGES - WHATSOEVER RESULTING FROM LOSS OF USE, DATA OR PROFITS, WHETHER IN AN - ACTION OF CONTRACT, NEGLIGENCE OR OTHER TORTIOUS ACTION, ARISING OUT OF - OR IN CONNECTION WITH THE USE OR PERFORMANCE OF THIS SOFTWARE. - - */ - -#ifndef MJSTATE_DATA_H -#define MJSTATE_DATA_H - -#include -#include - -#include -#include "simulation/dynamics/_GeneralModuleFiles/stateData.h" - -/** - * @brief Class representing the entire joint state data in a MuJoCo simulation. - * - * This class manages the position states (`qpos`) and their derivatives in a - * MuJoCo simulation. It extends `StateData` to include specific - * configurations and propagation logic for position states. - */ -class MJQPosStateData : public StateData -{ -public: - /** - * @brief Constructs an MJQPosStateData object. - * - * @param inName The name of the state. - * @param newState The initial state matrix. - */ - MJQPosStateData(std::string inName, const Eigen::MatrixXd& newState) - : StateData(std::move(inName), newState){}; - - /** - * @brief Creates a clone of the current state data object. - * - * @return A unique pointer to the cloned `StateData` object. - */ - virtual std::unique_ptr clone() const override; - - /** - * @brief Configures the state data with the given MuJoCo model. - * - * This method initializes the state and derivative sizes - * based on the given MuJoCo model. - * - * @param mujocoModel Pointer to the MuJoCo model. - */ - void configure(mjModel* mujocoModel); - - /** - * @brief Propagates the state over a time step. - * - * This method integrates the position state using the state derivative - * over the given time step:: - * - * \f[ - * x \mathrel{+}= f(t,x)\,h + g_0(t,x)\,\mathrm{pseudoStep}[0] + g_1(t,x)\,\mathrm{pseudoStep}[1] + \cdots - * \f] - * - * @param h The time step for propagation. - * @param pseudoStep For states driven by stochastic dynamics, this - * represents the random pseudotimestep. The length of this input must - * match the number of noise sources of this state (``getNumNoiseSources()``) - */ - void propagateState(double h, std::vector pseudoStep = {}) override; - -protected: - mjModel* mujocoModel; ///< Pointer to the MuJoCo model associated with the state. -}; - -#endif diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp index 4309d93003a..eed25533f23 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp @@ -95,20 +95,11 @@ void MJScene::Reset(uint64_t CurrentSimNanos) void MJScene::initializeDynamics() { - // Split MuJoCo's qpos/qvel into four states so the adaptive integrator - // can scale tolerances independently for free-joint translation - // (orbital metres / metres-per-second) and free-joint rotation + - // hinges/slides (radians / radians-per-second). A single combined - // state mixes those scales: the orbital component dominates the state - // norm and forces the relative-tolerance band so wide that errors in - // the stiff hinge ODE pass through unflagged, which is what produced - // the spurious panel spin we hit earlier. - this->qposTranslationState = this->dynManager.registerState(1, 1, "mujocoQposTranslation"); - this->qposAttitudeState = this->dynManager.registerState( - 1, 1, "mujocoQposAttitude"); - this->qvelTranslationState = this->dynManager.registerState(1, 1, "mujocoQvelTranslation"); - this->qvelAttitudeState = this->dynManager.registerState(1, 1, "mujocoQvelAttitude"); - this->actState = this->dynManager.registerState(1, 1, "mujocoAct"); + // Each joint owns its own qpos/qvel StateData(s); registering them + // per-body means the adaptive integrator scales tolerances independently + // per joint, at each joint's natural scale. The actuator state is the + // only scene-wide state that remains a bulk array. + this->actState = this->dynManager.registerState(1, 1, "mujocoAct"); for (auto&& body : this->spec.getBodies()) { body.registerStates(DynParamRegisterer( @@ -125,10 +116,16 @@ void MJScene::initializeDynamics() this->spec.configure(); } - // Now that the MuJoCo model is compiled, joint addresses are known - // and we can build the bidirectional maps between full mjData indices - // and split-state indices. - this->rebuildStateIndexMaps(); + // After spec compile the joint addresses (qposAdr/qvelAdr) are known and + // mjData::qpos / mjData::qvel hold the values declared in the XML. Seed + // each joint's StateData from those values so user-visible state matches + // the model out of the box. + { + auto d = this->spec.getMujocoData(); + for (auto&& body : this->spec.getBodies()) { + body.seedJointStatesFromMujoco(d); + } + } // Register the states of the models in the dynamics task std::unordered_set alreadyRegisteredModels; @@ -237,56 +234,15 @@ void MJScene::equationsOfMotion(double t, double timeStep) "s in MJScene with ID: " + std::to_string(moduleID)); } - // Replace the floating-frame translational acceleration for each free joint. - // In the co-moving frame, MuJoCo solves the dynamics with the body at the - // origin, so the qacc translation it produces is just floating-point noise - // from the co-moving cancellation — keeping it would drive the RKF45 - // step-size controller off the attitude PID timescale. We therefore - // overwrite it with the true inertial translational acceleration: the sum - // of gravity accelerations from all of the body's gravity sources evaluated - // at the body's saved (un-zeroed) inertial position, or zero if the body - // has no gravity sources (the original zero-gravity behaviour). - for (auto& sf : savedFrames) { - int i = sf.fj->getQvelAdr(); - Eigen::Vector3d gravAccel = sf.fj->getBody().computeGravityAt(sf.r); - mujocoData->qacc[i] = gravAccel[0]; - mujocoData->qacc[i + 1] = gravAccel[1]; - mujocoData->qacc[i + 2] = gravAccel[2]; - } - - // Push computed velocities and accelerations into the split-state - // derivatives. For qpos states the time derivative is the corresponding - // qvel state (which the attitude state interprets as angular velocity - // for quaternion blocks); for qvel states it is the qacc just computed. + // Each joint pulls its derivatives from the freshly computed mjData. + // Quaternion-flavoured joints (ball, free-attitude) set their qpos + // derivative to body angular velocity (consumed by QuaternionStateData), + // every other DOF uses the standard qpos' = qvel / qvel' = qacc. { - Eigen::VectorXd dQposTran(this->qposTranslationIdx.size()); - for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { - // Translation qpos derivative == translation qvel. - dQposTran[k] = mujocoData->qvel[this->qvelTranslationIdx[k]]; - } - this->qposTranslationState->setDerivative(dQposTran); - - Eigen::VectorXd dQposAtt(this->qvelAttitudeIdx.size()); - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - // Attitude qpos stateDeriv is sized as qvel-space (3 angular - // velocities per quaternion + 1 per linear DOF); MJQPosAttitudeStateData - // consumes it via mju_quatIntegrate (for quaternion blocks) and - // Euler step (for linear DOFs). - dQposAtt[k] = mujocoData->qvel[this->qvelAttitudeIdx[k]]; - } - this->qposAttitudeState->setDerivative(dQposAtt); - - Eigen::VectorXd dQvelTran(this->qvelTranslationIdx.size()); - for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { - dQvelTran[k] = qacc[this->qvelTranslationIdx[k]]; - } - this->qvelTranslationState->setDerivative(dQvelTran); - - Eigen::VectorXd dQvelAtt(this->qvelAttitudeIdx.size()); - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - dQvelAtt[k] = qacc[this->qvelAttitudeIdx[k]]; + auto d = this->spec.getMujocoData(); + for (auto&& body : this->spec.getBodies()) { + body.setJointDerivativesFromMujoco(d); } - this->qvelAttitudeState->setDerivative(dQvelAtt); } // Also copy the derivative of the actuator states, if we have them @@ -358,42 +314,6 @@ void MJScene::saveToFile(std::string filename) } } -StateData* -MJScene::getQposTranslationState() -{ - if (!this->qposTranslationState) { - logAndThrow("Tried to get qpos translation state before initialization."); - } - return this->qposTranslationState; -} - -MJQPosAttitudeStateData* -MJScene::getQposAttitudeState() -{ - if (!this->qposAttitudeState) { - logAndThrow("Tried to get qpos attitude state before initialization."); - } - return this->qposAttitudeState; -} - -StateData* -MJScene::getQvelTranslationState() -{ - if (!this->qvelTranslationState) { - logAndThrow("Tried to get qvel translation state before initialization."); - } - return this->qvelTranslationState; -} - -StateData* -MJScene::getQvelAttitudeState() -{ - if (!this->qvelAttitudeState) { - logAndThrow("Tried to get qvel attitude state before initialization."); - } - return this->qvelAttitudeState; -} - StateData* MJScene::getActState() { if (!this->actState) { @@ -402,42 +322,6 @@ StateData* MJScene::getActState() return this->actState; } -double& -MJScene::qposRef(int fullQposIdx) -{ - int splitIdx = this->fullQposToSplitIdx.at(fullQposIdx); - return this->fullQposIsTranslation.at(fullQposIdx) - ? this->qposTranslationState->state(splitIdx) - : this->qposAttitudeState->state(splitIdx); -} - -double -MJScene::qposGet(int fullQposIdx) const -{ - int splitIdx = this->fullQposToSplitIdx.at(fullQposIdx); - return this->fullQposIsTranslation.at(fullQposIdx) - ? this->qposTranslationState->state(splitIdx) - : this->qposAttitudeState->state(splitIdx); -} - -double& -MJScene::qvelRef(int fullQvelIdx) -{ - int splitIdx = this->fullQvelToSplitIdx.at(fullQvelIdx); - return this->fullQvelIsTranslation.at(fullQvelIdx) - ? this->qvelTranslationState->state(splitIdx) - : this->qvelAttitudeState->state(splitIdx); -} - -double -MJScene::qvelGet(int fullQvelIdx) const -{ - int splitIdx = this->fullQvelToSplitIdx.at(fullQvelIdx); - return this->fullQvelIsTranslation.at(fullQvelIdx) - ? this->qvelTranslationState->state(splitIdx) - : this->qvelAttitudeState->state(splitIdx); -} - void MJScene::printMujocoModelDebugInfo(const std::string& path) { @@ -560,26 +444,11 @@ MJScene::addForceTorqueActuator(const std::string& name, const MJSite& site) void MJScene::updateMujocoArraysFromStates() { auto mujocoModel = this->getMujocoModel(); - auto mujocoData = this->getMujocoData(); - - // Reassemble mjData::qpos from the two split position states. - for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { - mujocoData->qpos[this->qposTranslationIdx[k]] = - this->qposTranslationState->state(k); - } - for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { - mujocoData->qpos[this->qposAttitudeIdx[k]] = - this->qposAttitudeState->state(k); - } + auto mujocoData = this->getMujocoData(); - // Reassemble mjData::qvel from the two split velocity states. - for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { - mujocoData->qvel[this->qvelTranslationIdx[k]] = - this->qvelTranslationState->state(k); - } - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - mujocoData->qvel[this->qvelAttitudeIdx[k]] = - this->qvelAttitudeState->state(k); + // Each joint copies its owned state into mjData at its qposAdr/qvelAdr. + for (auto&& body : this->spec.getBodies()) { + body.syncJointStatesToMujoco(mujocoData); } if (mujocoModel->na > 0) { @@ -592,29 +461,21 @@ void MJScene::updateMujocoArraysFromStates() Eigen::VectorXd MJScene::assembleFullQpos() { + // Sync joints into mjData first (their state may have advanced past the + // last sync), then read straight from the contiguous mjData::qpos. + this->updateMujocoArraysFromStates(); auto m = this->spec.getMujocoModel(); - Eigen::VectorXd qpos = Eigen::VectorXd::Zero(m->nq); - for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { - qpos(this->qposTranslationIdx[k]) = this->qposTranslationState->state(k); - } - for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { - qpos(this->qposAttitudeIdx[k]) = this->qposAttitudeState->state(k); - } - return qpos; + auto d = this->spec.getMujocoData(); + return Eigen::Map(d->qpos, m->nq); } Eigen::VectorXd MJScene::assembleFullQvel() { + this->updateMujocoArraysFromStates(); auto m = this->spec.getMujocoModel(); - Eigen::VectorXd qvel = Eigen::VectorXd::Zero(m->nv); - for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { - qvel(this->qvelTranslationIdx[k]) = this->qvelTranslationState->state(k); - } - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - qvel(this->qvelAttitudeIdx[k]) = this->qvelAttitudeState->state(k); - } - return qvel; + auto d = this->spec.getMujocoData(); + return Eigen::Map(d->qvel, m->nv); } void @@ -626,119 +487,3 @@ MJScene::writeOutputStateMessages(uint64_t CurrentSimNanos) stateOutMsg.write(&stateOutMsgPayload, this->moduleID, CurrentSimNanos); } - -void -MJScene::rebuildStateIndexMaps() -{ - auto m = this->spec.getMujocoModel(); - - // Categorise joints by type using MuJoCo's compiled tables (instead of - // walking the MJBody hierarchy) so we get joint addresses in the same - // order MuJoCo lays them out in qpos/qvel. - std::vector freeJointIds; - std::vector ballJointIds; - std::vector scalarJointIds; - for (int j = 0; j < m->njnt; ++j) { - switch (m->jnt_type[j]) { - case mjJNT_FREE: freeJointIds.push_back(j); break; - case mjJNT_BALL: ballJointIds.push_back(j); break; - case mjJNT_HINGE: - case mjJNT_SLIDE: scalarJointIds.push_back(j); break; - } - } - - this->qposTranslationIdx.clear(); - this->qvelTranslationIdx.clear(); - this->qposAttitudeIdx.clear(); - this->qvelAttitudeIdx.clear(); - - // Translation state: free-joint translational DOFs only. - for (int jid : freeJointIds) { - int qposAdr = m->jnt_qposadr[jid]; - int qvelAdr = m->jnt_dofadr[jid]; - for (int k = 0; k < 3; ++k) { - this->qposTranslationIdx.push_back(qposAdr + k); - this->qvelTranslationIdx.push_back(qvelAdr + k); - } - } - - // Attitude state — quaternion blocks first (free-joint rotation, then - // ball joints), then scalar linear DOFs. This ordering matches - // MJQPosAttitudeStateData::propagateState, which expects all - // quaternions packed together at the front. - for (int jid : freeJointIds) { - int qposAdr = m->jnt_qposadr[jid]; - int qvelAdr = m->jnt_dofadr[jid]; - for (int k = 0; k < 4; ++k) this->qposAttitudeIdx.push_back(qposAdr + 3 + k); - for (int k = 0; k < 3; ++k) this->qvelAttitudeIdx.push_back(qvelAdr + 3 + k); - } - for (int jid : ballJointIds) { - int qposAdr = m->jnt_qposadr[jid]; - int qvelAdr = m->jnt_dofadr[jid]; - for (int k = 0; k < 4; ++k) this->qposAttitudeIdx.push_back(qposAdr + k); - for (int k = 0; k < 3; ++k) this->qvelAttitudeIdx.push_back(qvelAdr + k); - } - for (int jid : scalarJointIds) { - this->qposAttitudeIdx.push_back(m->jnt_qposadr[jid]); - this->qvelAttitudeIdx.push_back(m->jnt_dofadr[jid]); - } - - int numQuaternions = static_cast(freeJointIds.size() + ballJointIds.size()); - int numLinearDofs = static_cast(scalarJointIds.size()); - - // Size the split states. qposAttitudeState has separate state/deriv - // sizes (4 vs 3 per quaternion); the others use the default StateData - // layout where state and stateDeriv are the same size. - this->qposAttitudeState->configure(numQuaternions, numLinearDofs); - - auto resizeMatching = [](StateData* s, int n) { - s->state.resize(n, 1); - s->stateDeriv.resize(n, 1); - }; - resizeMatching(this->qposTranslationState, static_cast(this->qposTranslationIdx.size())); - resizeMatching(this->qvelTranslationState, static_cast(this->qvelTranslationIdx.size())); - resizeMatching(this->qvelAttitudeState, static_cast(this->qvelAttitudeIdx.size())); - - // Reverse maps for joint-side accessors (qposRef/qvelRef). - this->fullQposToSplitIdx.assign(m->nq, -1); - this->fullQposIsTranslation.assign(m->nq, false); - for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { - int fullIdx = this->qposTranslationIdx[k]; - this->fullQposToSplitIdx[fullIdx] = static_cast(k); - this->fullQposIsTranslation[fullIdx] = true; - } - for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { - int fullIdx = this->qposAttitudeIdx[k]; - this->fullQposToSplitIdx[fullIdx] = static_cast(k); - this->fullQposIsTranslation[fullIdx] = false; - } - this->fullQvelToSplitIdx.assign(m->nv, -1); - this->fullQvelIsTranslation.assign(m->nv, false); - for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { - int fullIdx = this->qvelTranslationIdx[k]; - this->fullQvelToSplitIdx[fullIdx] = static_cast(k); - this->fullQvelIsTranslation[fullIdx] = true; - } - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - int fullIdx = this->qvelAttitudeIdx[k]; - this->fullQvelToSplitIdx[fullIdx] = static_cast(k); - this->fullQvelIsTranslation[fullIdx] = false; - } - - // Seed the split states from the freshly compiled mjData so the user - // sees the qpos/qvel values declared in the XML before they make any - // setPosition/setVelocity calls. - auto d = this->spec.getMujocoData(); - for (size_t k = 0; k < this->qposTranslationIdx.size(); ++k) { - this->qposTranslationState->state(k) = d->qpos[this->qposTranslationIdx[k]]; - } - for (size_t k = 0; k < this->qposAttitudeIdx.size(); ++k) { - this->qposAttitudeState->state(k) = d->qpos[this->qposAttitudeIdx[k]]; - } - for (size_t k = 0; k < this->qvelTranslationIdx.size(); ++k) { - this->qvelTranslationState->state(k) = d->qvel[this->qvelTranslationIdx[k]]; - } - for (size_t k = 0; k < this->qvelAttitudeIdx.size(); ++k) { - this->qvelAttitudeState->state(k) = d->qvel[this->qvelAttitudeIdx[k]]; - } -} diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h index 516d525806c..6d2fb614ada 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -28,7 +28,6 @@ #include "MJUtils.h" #include "architecture/_GeneralModuleFiles/sys_model_task.h" #include "architecture/msgPayloadDefCpp/MJSceneStateMsgPayload.h" -#include "MJQPosAttitudeStateData.h" #include "simulation/dynamics/_GeneralModuleFiles/dynamicObject.h" #include @@ -568,55 +567,6 @@ class MJScene : public DynamicObject */ mjData* getMujocoData() { return this->spec.getMujocoData(); } - /** - * @brief Retrieves the free-joint translational position state data. - * - * Holds the 3-DOF translational positions of every free joint, packed - * contiguously. Lives in its own state (separate from - * `getQposAttitudeState`) so the adaptive integrator can scale tolerances - * for orbital position separately from radian-scale attitude/hinge DOFs. - * - * @return Pointer to the translational `StateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - StateData* getQposTranslationState(); - - /** - * @brief Retrieves the attitude/hinges position state data. - * - * Holds free-joint quaternions, ball-joint quaternions, and all scalar - * (hinge/slide) joint positions. Quaternion blocks are integrated with - * `mju_quatIntegrate`; linear DOFs use the default Euler step. - * - * @return Pointer to the attitude `MJQPosAttitudeStateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - MJQPosAttitudeStateData* getQposAttitudeState(); - - /** - * @brief Retrieves the free-joint translational velocity state data. - * - * Holds the 3-DOF translational velocities of every free joint, packed - * contiguously, separate from attitude/hinge velocities for tolerance - * scaling. - * - * @return Pointer to the translational velocity `StateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - StateData* getQvelTranslationState(); - - /** - * @brief Retrieves the attitude/hinges velocity state data. - * - * Holds free-joint angular velocities, ball-joint angular velocities, - * and all scalar joint velocities. It is the integrator-side derivative - * counterpart of `getQposAttitudeState`. - * - * @return Pointer to the attitude velocity `StateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - StateData* getQvelAttitudeState(); - /** * @brief Retrieves the actuation state data. * @@ -625,26 +575,6 @@ class MJScene : public DynamicObject */ StateData* getActState(); - /** - * @brief Returns a reference to the qpos element at the given full - * (mjData::qpos) index, dispatching to the translation or attitude split - * state. Used by joint setters to mutate state by full mujoco-data - * address without knowing which split state owns the element. - * - * @throw std::out_of_range if the index has not been mapped (caller - * passed an invalid index or `initializeDynamics` did not run). - */ - double& qposRef(int fullQposIdx); - - /** Const-read variant of `qposRef`. */ - double qposGet(int fullQposIdx) const; - - /** Same as `qposRef` but for the qvel/qacc indexing space. */ - double& qvelRef(int fullQvelIdx); - - /** Const-read variant of `qvelRef`. */ - double qvelGet(int fullQvelIdx) const; - /** * @brief Prints MuJoCo model debug information to a file. * @@ -703,14 +633,6 @@ class MJScene : public DynamicObject */ void writeOutputStateMessages(uint64_t CurrentSimNanos); - /** - * @brief Builds the index maps that translate full mjData qpos/qvel - * indices to/from indices in the split state vectors. Run once per - * `initializeDynamics`, after the MuJoCo spec has been compiled and - * joint addresses are known. - */ - void rebuildStateIndexMaps(); - protected: MJSpec spec; ///< `MJSpec` (MuJoCo model specification wrapper) associated with this scene. bool mjModelConstStale = false; ///< Flag indicating stale model constants. @@ -720,29 +642,7 @@ class MJScene : public DynamicObject SysModelTask dynamicsDiffusionTask; ///< Task managing models involved in the diffusion stochastic dynamics of this scene. std::vector> ownedSysModel; ///< System models that should be cleared on this scene destruction. - StateData* qposTranslationState = nullptr; ///< Free-joint translational positions (3 per free joint). - MJQPosAttitudeStateData* qposAttitudeState = nullptr; ///< Quaternions (free + ball joints) and scalar joint positions. - StateData* qvelTranslationState = nullptr; ///< Free-joint translational velocities (3 per free joint). - StateData* qvelAttitudeState = nullptr; ///< Angular velocities (free + ball) and scalar joint velocities. StateData* actState = nullptr; ///< Actuator state data. - - /** mjData::qpos index for each element of qposTranslationState. */ - std::vector qposTranslationIdx; - /** mjData::qvel index for each element of qvelTranslationState. */ - std::vector qvelTranslationIdx; - /** mjData::qpos index for each element of qposAttitudeState. */ - std::vector qposAttitudeIdx; - /** mjData::qvel index for each element of qvelAttitudeState. */ - std::vector qvelAttitudeIdx; - - /** Reverse map: full qpos index => index in its split state. */ - std::vector fullQposToSplitIdx; - /** Reverse map: full qpos index -> true if owned by translation state. */ - std::vector fullQposIsTranslation; - /** Reverse map for the qvel space. */ - std::vector fullQvelToSplitIdx; - /** Reverse map for the qvel space (true if translation). */ - std::vector fullQvelIsTranslation; }; template diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp index c624772da98..8ac7021276b 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp @@ -293,10 +293,10 @@ void MJSpec::configure() equality.configure(this->model.get()); } - // Sizing/seeding the qpos/qvel split states from mjData is done by - // MJScene::rebuildStateIndexMaps, which is invoked by - // MJScene::initializeDynamics after this configure() returns. Only the - // act state can be sized here because it's a flat copy with no split. + // qpos/qvel state ownership belongs to the individual joints; their + // states are seeded by MJScene::initializeDynamics after configure() + // returns (via MJBody::seedJointStatesFromMujoco). Only the actuator + // state is sized here because it's a flat scene-wide array. this->scene.getActState()->state.resize(this->model->na, 1); this->scene.getActState()->stateDeriv.resize(this->model->na, 1); std::copy_n(this->data->act, this->model->na, this->scene.getActState()->state.data());