diff --git a/docs/source/Support/bskReleaseNotesSnippets/mujoco-vizard.rst b/docs/source/Support/bskReleaseNotesSnippets/mujoco-vizard.rst new file mode 100644 index 00000000000..4faa139bb3c --- /dev/null +++ b/docs/source/Support/bskReleaseNotesSnippets/mujoco-vizard.rst @@ -0,0 +1,3 @@ +- Added ``getBodyNames()``, ``getBodyParentName()``, and ``getGeomInfos()`` methods to :ref:`MJScene` and ``MJSpec``, along with the new ``MJGeomInfo`` struct, to expose body hierarchy and geometry data to Python. +- Updated :ref:`vizSupport` ``enableUnityVisualization`` to accept :ref:`MJScene` objects. Automatically discover body hierarchy and generate Vizard models from MuJoCo data. +- Updated :ref:`scenarioDeployPanels` to support Vizard live stream visualization of the MuJoCo multi-body scene. diff --git a/examples/mujoco/scenarioDeployPanels.py b/examples/mujoco/scenarioDeployPanels.py index a491aba4fec..d410b049b84 100644 --- a/examples/mujoco/scenarioDeployPanels.py +++ b/examples/mujoco/scenarioDeployPanels.py @@ -67,7 +67,7 @@ from Basilisk.simulation import mujoco from Basilisk.utilities import SimulationBaseClass -from Basilisk.utilities import macros +from Basilisk.utilities import macros, vizSupport from Basilisk.simulation import StatefulSysModel from Basilisk.architecture import messaging from Basilisk.simulation import svIntegrators @@ -173,7 +173,7 @@ def generateProfiles( return interpolators -def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = False): +def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = False, vizard: bool = False): """Main function, see scenario description. Args: @@ -183,6 +183,8 @@ def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = Fa Defaults to False. visualize (bool, optional): If True, the ``MJScene`` visualization tool is run on the simulation results. Defaults to False. + vizard (bool, optional): If True, the simulation is visualized using Vizard. + Defaults to False. """ dt = 1 # s @@ -307,6 +309,9 @@ def run(initialSpin: bool = False, showPlots: bool = False, visualize: bool = Fa stateRecorder = scene.stateOutMsg.recorder() scSim.AddModelToTask("test", stateRecorder) + if vizard: + vizSupport.enableUnityVisualization(scSim, "test", scene, liveStream=True) + # Initialize the simulation and set the initial angles of the joints. # NOTE: the simulation MUST be initialized before setting the initial # joint angles and setting the initial attitude rate of the hub. @@ -447,4 +452,4 @@ def UpdateState(self, CurrentSimNanos: int): if __name__ == "__main__": - run(False, True, True) + run(False, True, True, vizard=True) diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp index 4c95fcfcbb9..eb9fed5cd4e 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.cpp @@ -329,6 +329,21 @@ void MJScene::printMujocoModelDebugInfo(const std::string& path) mj_printModel(this->getMujocoModel(), path.c_str()); } +std::vector MJScene::getBodyNames() const +{ + return this->spec.getBodyNames(); +} + +std::string MJScene::getBodyParentName(const std::string& bodyName) const +{ + return this->spec.getBodyParentName(bodyName); +} + +std::vector MJScene::getGeomInfos() const +{ + return this->spec.getGeomInfos(); +} + MJBody& MJScene::getBody(const std::string& name) { auto& bodies = this->spec.getBodies(); diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h index f9035338ad7..4e54a6c341b 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.h @@ -70,6 +70,39 @@ class MJScene : public DynamicObject MJScene& operator=(const MJScene&) = delete; MJScene& operator=(MJScene&&) = delete; + /** + * @brief Retrieves the names of all bodies in the scene. + * + * This method does not trigger model recompilation, so it + * can safely be called before initialization. + * + * @return A vector of body names. + */ + std::vector getBodyNames() const; + + /** + * @brief Retrieves the parent body name of a given body. + * + * Returns an empty string if the body's parent is the worldbody + * or the body is not found. This method does not trigger model + * recompilation, so it can safely be called before initialization. + * + * @param bodyName The name of the body. + * @return The parent body's name, or empty string if none. + */ + std::string getBodyParentName(const std::string& bodyName) const; + + /** + * @brief Retrieves geometry information for all geoms in the scene. + * + * Each entry describes a geom's shape, size, position, orientation, + * and color, along with which body it belongs to. This method does + * not trigger model recompilation. + * + * @return A vector of `MJGeomInfo` structs. + */ + std::vector getGeomInfos() const; + /** * @brief Retrieves a body by name. * diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.swg b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.swg index 56ff6128fe9..94ab2363447 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.swg +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJScene.swg @@ -28,6 +28,17 @@ from typing import List %include "std_vector.i" %include "std_string.i" %template() std::vector; +%template() std::vector; + +struct MJGeomInfo { + std::string bodyName; + int type; + std::vector size; + std::vector pos; + std::vector quat; + std::vector rgba; +}; +%template(MJGeomInfoVector) std::vector; %include "simulation/dynamics/_GeneralModuleFiles/dynParamManager.i" %include "simulation/dynamics/_GeneralModuleFiles/dynamicObject.h" diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp index e013c5b6946..4a4bccb15ba 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.cpp @@ -207,6 +207,53 @@ mjData* MJSpec::getMujocoData() return this->data.get(); } +std::vector MJSpec::getBodyNames() const +{ + std::vector names; + for (auto&& body : this->bodies) { + names.push_back(body.getName()); + } + return names; +} + +std::string MJSpec::getBodyParentName(const std::string& bodyName) const +{ + // Use the existing model directly without recompile since the body structure + // should not change after initial compilation + int bodyId = mj_name2id(this->model.get(), mjOBJ_BODY, bodyName.c_str()); + if (bodyId <= 0) return ""; + int parentId = this->model->body_parentid[bodyId]; + if (parentId == 0) return ""; + return std::string(this->model->names + this->model->name_bodyadr[parentId]); +} + +std::vector MJSpec::getGeomInfos() const +{ + std::vector geoms; + auto m = this->model.get(); + + for (int i = 0; i < m->ngeom; i++) { + MJGeomInfo info; + + int bodyId = m->geom_bodyid[i]; + if (bodyId == 0) continue; + info.bodyName = std::string(m->names + m->name_bodyadr[bodyId]); + + info.type = m->geom_type[i]; + info.size = {m->geom_size[i * 3], m->geom_size[i * 3 + 1], m->geom_size[i * 3 + 2]}; + info.pos = {m->geom_pos[i * 3], m->geom_pos[i * 3 + 1], m->geom_pos[i * 3 + 2]}; + info.quat = {m->geom_quat[i * 4], m->geom_quat[i * 4 + 1], + m->geom_quat[i * 4 + 2], m->geom_quat[i * 4 + 3]}; + info.rgba = {static_cast(m->geom_rgba[i * 4]), + static_cast(m->geom_rgba[i * 4 + 1]), + static_cast(m->geom_rgba[i * 4 + 2]), + static_cast(m->geom_rgba[i * 4 + 3])}; + + geoms.push_back(std::move(info)); + } + return geoms; +} + bool MJSpec::hasActuator(const std::string& name) { return std::find_if(std::begin(actuators), std::end(actuators), [&](auto&& obj) { diff --git a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.h b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.h index b57e4c67a88..b22a71b193d 100644 --- a/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.h +++ b/src/simulation/mujocoDynamics/_GeneralModuleFiles/MJSpec.h @@ -35,6 +35,19 @@ namespace Eigen using Vector6d = Eigen::Matrix; } +/** + * @brief Describes the geometry (shape, size, position, orientation, color) of a single + * MuJoCo geom, suitable for visualization. + */ +struct MJGeomInfo { + std::string bodyName; ///< Name of the body this geom belongs to. + int type; ///< MuJoCo geom type (e.g. mjGEOM_BOX=6, mjGEOM_SPHERE=2, mjGEOM_CYLINDER=5, etc.). + std::vector size; ///< Size parameters (3 elements). + std::vector pos; ///< Position in body frame (3 elements). + std::vector quat; ///< Quaternion in body frame (4 elements: w, x, y, z). + std::vector rgba; ///< Color (4 elements, 0-1 range). +}; + class MJScene; /** @@ -164,6 +177,39 @@ class MJSpec */ std::list& getEqualities() { return equalities; } + /** + * @brief Retrieves the names of all bodies in the specification. + * + * This method does not trigger model recompilation, so it + * can be called before initialization. + * + * @return A vector of body names. + */ + std::vector getBodyNames() const; + + /** + * @brief Retrieves the parent body name of a given body. + * + * Uses the existing compiled model (without triggering recompilation) + * to look up the parent relationship. Returns an empty string if + * the body's parent is the worldbody or the body is not found. + * + * @param bodyName The name of the body. + * @return The parent body's name, or empty string if none. + */ + std::string getBodyParentName(const std::string& bodyName) const; + + /** + * @brief Retrieves geometry information for all geoms in the model. + * + * Uses the existing compiled model (without triggering recompilation). + * Each entry describes a geom's shape, size, position, orientation, + * and color, along with which body it belongs to. + * + * @return A vector of `MJGeomInfo` structs. + */ + std::vector getGeomInfos() const; + /** * @brief Checks if an actuator exists with the given name. * diff --git a/src/tests/data/hypso1.tle b/src/tests/data/hypso1.tle new file mode 100644 index 00000000000..a2820cd6338 --- /dev/null +++ b/src/tests/data/hypso1.tle @@ -0,0 +1,3 @@ +HYPSO 1 +1 51053U 22002BX 25279.47924866 .00028852 00000-0 56053-3 0 9998 +2 51053 97.3197 349.0390 0004257 136.0807 224.0782 15.47578936207406 diff --git a/src/tests/data/oneWeb25.tle b/src/tests/data/oneWeb25.tle new file mode 100644 index 00000000000..fa3c96ecaa5 --- /dev/null +++ b/src/tests/data/oneWeb25.tle @@ -0,0 +1,18 @@ +ONEWEB-0012 +1 44057U 19010A 25275.66668980 -.00044498 00000+0 -11697+0 0 9994 +2 44057 87.8965 280.1128 0001750 78.6519 126.1333 13.16614029 5036 +ONEWEB-0010 +1 44058U 19010B 25275.57281529 .00000022 00000+0 23658-4 0 9993 +2 44058 87.8967 280.1151 0001764 102.8888 257.2440 13.16593240317716 +ONEWEB-0008 +1 44059U 19010C 25275.44615373 .00000099 00000+0 22482-3 0 9995 +2 44059 87.8965 280.1373 0002086 88.2864 271.8506 13.16593437317814 +ONEWEB-0007 +1 44060U 19010D 25275.58075372 -.00000270 00000+0 -75002-3 0 9990 +2 44060 87.8998 310.6396 0001581 99.4188 260.7121 13.15547521317506 +ONEWEB-0006 +1 44061U 19010E 25275.66076213 -.00000298 00000+0 -82440-3 0 9999 +2 44061 87.9003 310.5649 0001714 101.6321 258.5002 13.15547636317564 +ONEWEB-0011 +1 44062U 19010F 25275.62907667 .00000180 00000+0 44426-3 0 9992 +2 44062 87.8998 310.5840 0002029 98.4046 261.7314 13.15547351317552 diff --git a/src/tests/data/spacestations.2le b/src/tests/data/spacestations.2le new file mode 100644 index 00000000000..abb7d26d283 --- /dev/null +++ b/src/tests/data/spacestations.2le @@ -0,0 +1,6 @@ +1 25544U 98067A 25275.08016394 .00015216 00000+0 27786-3 0 9999 +2 25544 51.6326 137.2684 0001066 187.1210 172.9763 15.49606939531746 +1 48274U 21035A 25274.88825115 .00047308 00000+0 55227-3 0 9999 +2 48274 41.4640 93.5081 0001842 163.6733 196.4166 15.60798035252800 +1 49044U 21066A 25274.82220399 .00016258 00000+0 29639-3 0 9994 +2 49044 51.6327 138.5459 0001045 186.4676 173.6300 15.49599116230443 diff --git a/src/utilities/vizSupport.py b/src/utilities/vizSupport.py index c8b3a229881..e71b0167f02 100644 --- a/src/utilities/vizSupport.py +++ b/src/utilities/vizSupport.py @@ -37,6 +37,13 @@ except ImportError: vizFound = False +try: + from Basilisk.simulation import mujoco + + mujocoFound = True +except ImportError: + mujocoFound = False + pauseFlag = False endFlag = False @@ -528,6 +535,102 @@ def updateTargetLineList(viz): return +def _createCustomModelsFromMJScene(viz, scene): + """Auto-create Vizard custom models from MuJoCo geom data. + + Maps MuJoCo geom types to Vizard primitive shapes and creates + custom models with the correct size, position, orientation, and color + for each body in the scene. + """ + import math + + # MuJoCo geom type constants (from mjtGeom enum) + _MJGEOM_PLANE = 0 + _MJGEOM_HFIELD = 1 + _MJGEOM_SPHERE = 2 + _MJGEOM_CAPSULE = 3 + _MJGEOM_ELLIPSOID = 4 + _MJGEOM_CYLINDER = 5 + _MJGEOM_BOX = 6 + _MJGEOM_MESH = 7 + + # Map MuJoCo geom types to Vizard primitive model paths + _GEOM_TYPE_MAP = { + _MJGEOM_BOX: "CUBE", + _MJGEOM_SPHERE: "SPHERE", + _MJGEOM_CYLINDER: "CYLINDER", + _MJGEOM_CAPSULE: "CYLINDER", + _MJGEOM_ELLIPSOID: "SPHERE", + } + + geomInfos = scene.getGeomInfos() + + # Group geoms by body name + bodyGeoms = {} + for geom in geomInfos: + bodyGeoms.setdefault(geom.bodyName, []).append(geom) + + for bodyName, geoms in bodyGeoms.items(): + for geom in geoms: + modelPath = _GEOM_TYPE_MAP.get(geom.type) + if modelPath is None: + continue + + # Compute scale from MuJoCo size + size = list(geom.size) + if geom.type == _MJGEOM_BOX: + # MuJoCo size = half-extents, Vizard CUBE default = 1x1x1 + scale = [2 * size[0], 2 * size[1], 2 * size[2]] + elif geom.type == _MJGEOM_SPHERE: + # MuJoCo size[0] = radius, Vizard SPHERE default diameter = 1 + scale = [2 * size[0], 2 * size[0], 2 * size[0]] + elif geom.type == _MJGEOM_CYLINDER: + # MuJoCo size[0] = radius, size[1] = half-height + # Vizard CYLINDER default = diameter 1, height 1 + scale = [2 * size[0], 2 * size[1], 2 * size[0]] + elif geom.type == _MJGEOM_CAPSULE: + # Approximate as cylinder + scale = [2 * size[0], 2 * size[1], 2 * size[0]] + elif geom.type == _MJGEOM_ELLIPSOID: + # MuJoCo size = semi-axes + scale = [2 * size[0], 2 * size[1], 2 * size[2]] + else: + scale = [1, 1, 1] + + # Position offset in body frame + offset = list(geom.pos) + + # Convert quaternion (w, x, y, z) to 3-2-1 Euler angles (z, y, x) + qw, qx, qy, qz = geom.quat + # Roll (x-axis rotation) + sinr_cosp = 2 * (qw * qx + qy * qz) + cosr_cosp = 1 - 2 * (qx * qx + qy * qy) + roll = math.atan2(sinr_cosp, cosr_cosp) + # Pitch (y-axis rotation) + sinp = 2 * (qw * qy - qz * qx) + sinp = max(-1, min(1, sinp)) + pitch = math.asin(sinp) + # Yaw (z-axis rotation) + siny_cosp = 2 * (qw * qz + qx * qy) + cosy_cosp = 1 - 2 * (qy * qy + qz * qz) + yaw = math.atan2(siny_cosp, cosy_cosp) + rotation = [yaw, pitch, roll] + + # Convert RGBA from 0-1 float to 0-255 int + rgba = list(geom.rgba) + color = [int(c * 255) for c in rgba] + + createCustomModel( + viz, + modelPath=modelPath, + simBodiesToModify=[bodyName], + scale=scale, + offset=offset, + rotation=rotation, + color=color, + ) + + customModelList = [] @@ -1108,7 +1211,10 @@ def enableUnityVisualization( simTaskName: task to which to add the vizInterface module scList: - :ref:`spacecraft` objects. Can be a single object or list of objects + :ref:`spacecraft` objects, :ref:`MJScene` objects, or tuples of (name, stateOutMsg). + Can be a single object or list of objects. + When an :ref:`MJScene` is provided, all bodies are auto-discovered and their + parent-child hierarchy is resolved from the MuJoCo model Keyword Args ------------ @@ -1267,6 +1373,52 @@ def enableUnityVisualization( spacecraftParentName = "" for sc in scList: + # Check if this is an MJScene object + if mujocoFound and isinstance(sc, mujoco.MJScene): + # Auto-discover all bodies their child hierarchy + bodyNames = sc.getBodyNames() + + # Find the hub body, set it as top-level spacecraft + hubName = None + for name in bodyNames: + if sc.getBody(name).isFree(): + hubName = name + break + if hubName is None: + raise ValueError( + "MJScene has no free body. Cannot determine the hub spacecraft." + ) + + # Add the hub body first as the top-level spacecraft + scData = vizInterface.VizSpacecraftData() + scData.spacecraftName = hubName + spacecraftParentName = hubName + scData.scStateInMsg.subscribeTo(sc.getBody(hubName).getOrigin().stateOutMsg) + scSim.vizMessenger.scData.push_back(scData) + + # Add all other bodies as child spacecraft, using the MuJoCo + # body hierarchy to resolve parentSpacecraftName + for name in bodyNames: + if name == hubName: + continue + scData = vizInterface.VizSpacecraftData() + scData.spacecraftName = name + + # Walk up the MuJoCo body tree to find the nearest ancestor that + # is in bodyNames + parentName = sc.getBodyParentName(name) + scData.parentSpacecraftName = parentName if parentName else hubName + + scData.scStateInMsg.subscribeTo(sc.getBody(name).getOrigin().stateOutMsg) + scSim.vizMessenger.scData.push_back(scData) + + # Auto-create custom models from MuJoCo geom data so Vizard renders + # the correct geometry for each body + _createCustomModelsFromMJScene(scSim.vizMessenger, sc) + + c += 1 + continue + # create spacecraft information container scData = vizInterface.VizSpacecraftData()