Add RobotMjcf

VirtualRobot/CMakeLists.txt

@@ -377,6 +377,7 @@ SET(SOURCES
     XML/mujoco/DummyMassBodySanitizer.cpp
     XML/mujoco/MergingBodySanitizer.cpp
     XML/mujoco/MujocoIO.cpp
+    XML/mujoco/RobotMjcf.cpp
 )
 
 SET(INCLUDES
@@ -600,6 +601,7 @@ SET(INCLUDES
     XML/mujoco/DummyMassBodySanitizer.h
     XML/mujoco/MergingBodySanitizer.h
     XML/mujoco/MujocoIO.h
+    XML/mujoco/RobotMjcf.h
     
 )
 

VirtualRobot/XML/mujoco/RobotMjcf.cpp

+#include "RobotMjcf.h"
+
+#include <boost/shared_ptr.hpp>
+
+#include <VirtualRobot/RobotNodeSet.h>
+
+#include <VirtualRobot/Nodes/RobotNodePrismatic.h>
+#include <VirtualRobot/Nodes/RobotNodeRevolute.h>
+
+
+
+namespace VirtualRobot::mujoco
+{
+
+RobotMjcf::RobotMjcf(RobotPtr robot) : robot(robot)
+{
+    document->setModelName(robot->getName());
+    
+    // Add robot root.
+    robotBody = document->worldbody().addBody(robot->getName(), robot->getName());
+    nodeBodies[robot->getName()] = robotBody;
+}
+
+mjcf::Document& RobotMjcf::getDocument()
+{
+    return *document;
+}
+
+const mjcf::Document& RobotMjcf::getDocument() const
+{
+    return *document;
+}
+
+RobotPtr RobotMjcf::getRobot() const
+{
+    return robot;
+}
+
+void RobotMjcf::setOutputFile(const std::filesystem::path& filePath)
+{
+    this->outputFile = filePath;
+}
+
+void RobotMjcf::addCompiler(const std::string& angle, bool balanceIneratia)
+{
+    document->compiler().angle = angle;
+    document->compiler().balanceinertia = balanceIneratia;
+}
+
+void RobotMjcf::addDefaultsClass(float meshScale)
+{
+    mjcf::DefaultClass defaultsClass = document->default_().addClass(robot->getName());
+    
+    defaultsClass.insertComment("Add default values for " + robot->getName() + " here.", true);
+    
+    mjcf::Joint joint = defaultsClass.getElement<mjcf::Joint>();
+    joint.frictionloss = 1;
+    joint.damping = 0;
+    
+    mjcf::Mesh mesh = defaultsClass.getElement<mjcf::Mesh>();
+    mesh.scale = Eigen::Vector3f::Constant(meshScale);
+    
+    mjcf::Geom geom = defaultsClass.getElement<mjcf::Geom>();
+    geom.condim = 4;
+    
+    mjcf::ActuatorPosition actPos = defaultsClass.getElement<mjcf::ActuatorPosition>();
+    actPos.kp = 1;
+    
+    mjcf::ActuatorVelocity actVel = defaultsClass.getElement<mjcf::ActuatorVelocity >();
+    actVel.kv = 1;
+}
+
+void RobotMjcf::addSkybox(const Eigen::Vector3f& rgb1, const Eigen::Vector3f& rgb2)
+{
+    document->asset().addSkyboxTexture(rgb1, rgb2);
+}
+
+mjcf::Body RobotMjcf::addNodeBody(RobotNodePtr node, mjcf::Body parent,
+                                  bool addJoint, bool addInertial)
+{
+    mjcf::Body body = parent.addBody(node->getName());
+
+    if (node->hasParent())
+    {
+        Eigen::Matrix4f pose = node->getTransformationFrom(node->getParent());
+        math::Helpers::ScaleTranslation(pose, lengthScale);
+        body.setPose(pose);
+    }
+    
+    if (addJoint && (node->isRotationalJoint() || node->isTranslationalJoint()))
+    {
+        addNodeJoint(node, body);
+    }
+    
+    if (addInertial)
+    {
+        addNodeInertial(node, body);
+    }
+    
+    return body;
+}
+
+mjcf::Joint RobotMjcf::addNodeJoint(RobotNodePtr node, mjcf::Body body)
+{
+    VR_ASSERT(node->isRotationalJoint() xor node->isTranslationalJoint());
+    
+    mjcf::Joint joint = body.addJoint();
+    joint.name = node->getName();
+    
+    // get the axis
+    Eigen::Vector3f axis;
+    if (node->isRotationalJoint())
+    {
+        RobotNodeRevolutePtr revolute = boost::dynamic_pointer_cast<RobotNodeRevolute>(node);
+        VR_ASSERT(revolute);
+        axis = revolute->getJointRotationAxisInJointCoordSystem();
+    }
+    else if (node->isTranslationalJoint())
+    {
+        RobotNodePrismaticPtr prismatic = boost::dynamic_pointer_cast<RobotNodePrismatic>(node);
+        VR_ASSERT(prismatic);
+        axis = prismatic->getJointTranslationDirectionJointCoordSystem();
+    }
+    
+    joint.type = node->isRotationalJoint() ? "hinge" : "slide";
+    joint.axis = axis;
+    joint.limited = !node->isLimitless();
+    
+    if (!node->isLimitless())
+    {
+        Eigen::Vector2f range = { node->getJointLimitLow(), node->getJointLimitHigh() };
+        if (node->isTranslationalJoint())
+        {
+            range *= lengthScale;
+        }
+        
+        // Mujoco does not like ranges where min >= max.
+        if (std::abs(range(0) - range(1)) < 1e-6f)
+        {
+            range(1) = range(0) + 1e-6f;
+        }
+        joint.range = range;
+    }
+    
+    return joint;
+}
+
+mjcf::Inertial RobotMjcf::addNodeInertial(RobotNodePtr node, mjcf::Body body)
+{
+    const Eigen::Matrix3f matrix = node->getInertiaMatrix();
+    if (matrix.isIdentity(document->getFloatCompPrecision()) 
+        && node->getMass() < document->getFloatCompPrecision())
+    {
+        // Dont set an inertial element and let it be derived automatically.
+        return { };
+    }
+    
+    mjcf::Inertial inertial = body.addInertial();
+    inertial.pos = node->getCoMLocal() * lengthScale;
+    inertial.mass = node->getMass() * massScale;
+    inertial.inertiaFromMatrix(matrix);
+    
+    return inertial;
+}
+
+mjcf::Body RobotMjcf::addNodeBody(RobotNodePtr node, bool addJoint, bool addInertial)
+{
+    // See whether body for the node already exists.
+    auto find = nodeBodies.find(node->getName());
+    if (find != nodeBodies.end())
+    {
+        // Exists => break recursion.
+        return find->second;
+    }
+    
+    // Check whether body exists for parent node.
+    mjcf::Body parent;
+    
+    if (node->getName() == robot->getRootNode()->getName())
+    {
+        parent = robotBody;
+    }
+    else
+    {
+        find = nodeBodies.find(node->getParent()->getName());
+        if (find != nodeBodies.end())
+        {
+            // Parent exists.
+            parent = find->second;
+        }
+        else
+        {
+            // Parent does not exist => create it first.
+            parent = addNodeBody(robot->getRobotNode(node->getParent()->getName()),
+                                 addJoint, addInertial);
+        }
+    }
+    
+    // Add body as child of parent.
+    mjcf::Body body = addNodeBody(node, parent, addJoint, addInertial);
+    nodeBodies[node->getName()] = body;
+
+    return body;
+}
+
+
+void RobotMjcf::addNodeBodies()
+{
+    addNodeBodies(robot->getRobotNodeNames());
+}
+
+void RobotMjcf::addNodeBodies(const RobotNodeSet& nodeSet)
+{
+    addNodeBodies(nodeSet.getNodeNames());
+}
+
+void RobotMjcf::addNodeBodies(const std::vector<std::string>& nodeNames)
+{
+    for (const std::string& nodeName : nodeNames)
+    {
+        addNodeBody(robot->getRobotNode(nodeName));
+    }
+}
+
+
+}

VirtualRobot/XML/mujoco/RobotMjcf.h

+#pragma once
+
+#include <filesystem>
+
+#include <VirtualRobot/Robot.h>
+#include <VirtualRobot/MJCF/Document.h>
+
+
+namespace VirtualRobot::mujoco
+{
+
+
+    class RobotMjcf
+    {
+    public:
+        
+        /// Construct with robot.
+        RobotMjcf(RobotPtr robot);
+        
+        
+        /// Get the document.
+        mjcf::Document& getDocument();
+        const mjcf::Document& getDocument() const;
+        
+        
+        // PREPARATION
+        
+        /// Get the robot.
+        RobotPtr getRobot() const;
+        
+        
+        
+        
+        /// Set the path to the output XML file.
+        void setOutputFile(const std::filesystem::path& filePath);
+        
+        
+        // MODIFIERS
+        
+        // META
+        
+        /// Make a compiler section.
+        void addCompiler(const std::string& angle = "radian", bool balanceIneratia = true);
+        
+        /// Make defaults class with robot's name.
+        void addDefaultsClass(float meshScale = 1.0f);
+        
+        
+        // ASSETS
+        
+        /// Add skybox texture.
+        void addSkybox(const Eigen::Vector3f& rgb1 = { .8f, .9f, .95f }, 
+                       const Eigen::Vector3f& rgb2 = { .4f, .6f, .80f });
+        
+        
+        // ROBOT NODES
+        
+        /// Add a body for the given robot node as child of `parent`.
+        mjcf::Body addNodeBody(RobotNodePtr node, mjcf::Body parent,
+                               bool addJoint = true, bool addInertial = true);
+        
+        /// Add a joint for the given node in `body` and return it.
+        mjcf::Joint addNodeJoint(RobotNodePtr node, mjcf::Body body);
+        /// Add an inertial for the given node in `body` and return it.
+        mjcf::Inertial addNodeInertial(RobotNodePtr node, mjcf::Body body);
+        
+        
+        /**
+         * @brief Add a body element for the given robot node.
+         * If its parent does not exist, create the parent body first.
+         * If `node` is the robot root node, it is attached to the robot body.
+         * @param addJoint Whether to add a joint, if applicable.
+         * @param addInertial Whether to add an inertial, if applicable.
+         * @return The node body.
+         */
+        mjcf::Body addNodeBody(RobotNodePtr node, bool addJoint = true, bool addInertial = true);
+
+        
+        /// Add bodies for all robot nodes.
+        void addNodeBodies();
+        /// Add bodies for the specified nodes.
+        void addNodeBodies(const RobotNodeSet& nodeSet);
+        /// Add bodies for the specified nodes.
+        void addNodeBodies(const std::vector<std::string>& nodeNames);
+        
+        
+    private:
+
+        /// The robot.
+        RobotPtr robot;
+        
+        /// The document. Using a pointer allows to easily reset it.
+        mjcf::DocumentPtr document { new mjcf::Document() };
+        
+        /// The path to the output XML file.
+        std::filesystem::path outputFile;
+        
+        
+        /// Scaling for lengths, such as positions and translations (to m).
+        float lengthScale = 1.f;
+        /// Scaling for lengths, such as positions and translations (to m).
+        float meshScale = 1.f;
+        /// Scaling for mass (to kg).
+        float massScale = 1.f;
+        
+        
+        /// The robot body. (Not part of the original robot structure.)
+        mjcf::Body robotBody;
+        
+        /// Map of robot node names to XML elements.
+        std::map<std::string, mjcf::Body> nodeBodies;
+        
+    };
+
+}