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 58c47907d14..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,35 +102,62 @@ 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(); + 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 = this->qposState->state(0); + 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 = this->qvelState->state(0); + 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(); + this->qposState->state(0) = value; + this->body.getSpec().getScene().markKinematicsAsStale(); } void MJScalarJoint::setVelocity(double value) { checkInitialized(); - - body.getSpec().getScene().getQvelState()->state(this->qvelAdr.value()) = value; - body.getSpec().getScene().markKinematicsAsStale(); + this->qvelState->state(0) = value; + this->body.getSpec().getScene().markKinematicsAsStale(); } MJSingleJointEquality @@ -135,49 +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& qposState = body.getSpec().getScene().getQposState()->state; - auto i = this->qposAdr.value(); - qposState.middleRows(i, 3) = position; +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); +} - body.getSpec().getScene().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& qvelState = body.getSpec().getScene().getQvelState()->state; - auto i = this->qvelAdr.value(); - qvelState.middleRows(i, 3) = velocity; +// --------------------------------------------------------------------------- +// MJFreeJoint +// --------------------------------------------------------------------------- - body.getSpec().getScene().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& qposState = body.getSpec().getScene().getQposState()->state; - 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); +} - qposState.middleRows(i + 3, 4) = Eigen::Vector4d{quat.w(), quat.x(), quat.y(), quat.z()}; - body.getSpec().getScene().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& qvelState = body.getSpec().getScene().getQvelState()->state; - auto i = this->qvelAdr.value(); - qvelState.middleRows(i + 3, 3) = attitudeRate; +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(); +} - body.getSpec().getScene().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/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 4c95fcfcbb9..eed25533f23 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp @@ -95,11 +95,10 @@ 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"); + // 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()) { @@ -117,6 +116,17 @@ void MJScene::initializeDynamics() this->spec.configure(); } + // 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; auto registerStatesOnSysModel = [this, &alreadyRegisteredModels](SysModel* sysModelPtr) @@ -224,12 +234,16 @@ 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()); - - // 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); + // 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. + { + auto d = this->spec.getMujocoData(); + for (auto&& body : this->spec.getBodies()) { + body.setJointDerivativesFromMujoco(d); + } + } // Also copy the derivative of the actuator states, if we have them if (this->spec.getMujocoModel()->na > 0) { @@ -300,31 +314,16 @@ void MJScene::saveToFile(std::string filename) } } -MJQPosStateData* MJScene::getQposState() -{ - if (!this->qposState) { - logAndThrow("Tried to get qpos state before initialization."); - } - return this->qposState; -} - -StateData* MJScene::getQvelState() -{ - if (!this->qvelState) { - logAndThrow("Tried to get qvel state before initialization."); - } - return this->qvelState; -} - 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) +void +MJScene::printMujocoModelDebugInfo(const std::string& path) { mj_printModel(this->getMujocoModel(), path.c_str()); } @@ -445,13 +444,12 @@ MJScene::addForceTorqueActuator(const std::string& name, const MJSite& site) void MJScene::updateMujocoArraysFromStates() { auto mujocoModel = this->getMujocoModel(); - auto mujocoData = this->getMujocoData(); + 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); - - std::copy_n(this->qvelState->state.data(), mujocoModel->nv, mujocoData->qvel); + // 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) { std::copy_n(this->actState->state.data(), mujocoModel->na, mujocoData->act); @@ -460,11 +458,32 @@ void MJScene::updateMujocoArraysFromStates() markKinematicsAsStale(); } -void MJScene::writeOutputStateMessages(uint64_t CurrentSimNanos) +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(); + auto d = this->spec.getMujocoData(); + return Eigen::Map(d->qpos, m->nq); +} + +Eigen::VectorXd +MJScene::assembleFullQvel() +{ + this->updateMujocoArraysFromStates(); + auto m = this->spec.getMujocoModel(); + auto d = this->spec.getMujocoData(); + return Eigen::Map(d->qvel, m->nv); +} + +void +MJScene::writeOutputStateMessages(uint64_t CurrentSimNanos) { - MJSceneStateMsgPayload stateOutMsgPayload{this->qposState->getState(), - this->qvelState->getState(), - this->actState->getState()}; + MJSceneStateMsgPayload stateOutMsgPayload{ this->assembleFullQpos(), + this->assembleFullQvel(), + this->actState->getState() }; stateOutMsg.write(&stateOutMsgPayload, this->moduleID, CurrentSimNanos); } diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h index f9035338ad7..6d2fb614ada 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -28,9 +28,10 @@ #include "MJUtils.h" #include "architecture/_GeneralModuleFiles/sys_model_task.h" #include "architecture/msgPayloadDefCpp/MJSceneStateMsgPayload.h" -#include "MJQPosStateData.h" #include "simulation/dynamics/_GeneralModuleFiles/dynamicObject.h" +#include + /** * @brief Represents a dynamic object that solves multi-body dynamics through MuJoCo. * @@ -566,22 +567,6 @@ class MJScene : public DynamicObject */ mjData* getMujocoData() { return this->spec.getMujocoData(); } - /** - * @brief Retrieves the position state data. - * - * @return Pointer to `MJQPosStateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - MJQPosStateData* getQposState(); - - /** - * @brief Retrieves the velocity state data. - * - * @return Pointer to `StateData`. - * @throw std::runtime_error If the state has not been initialized. - */ - StateData* getQvelState(); - /** * @brief Retrieves the actuation state data. * @@ -624,6 +609,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. @@ -646,9 +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. - MJQPosStateData* qposState; ///< Position state data. - StateData* qvelState; ///< Velocity state data. - StateData* actState; ///< Actuator state data. + StateData* actState = nullptr; ///< Actuator state data. }; template diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp index 7fa4e9a9669..8ac7021276b 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()); - + // 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());