fix: add MJCF robot methods to config and update examples

Author: senthurayyappan

examples/simulation/main.py

@@ -10,12 +10,13 @@
 import plotly
 from controllers import PIDController
 from loguru import logger
-from mods import modify_ballbot
+from mods import create_ballbot_config
 from mujoco.usd import exporter
 from scipy.spatial.transform import Rotation
 from transformations import compute_motor_torques
 
 from onshape_robotics_toolkit.connect import Client
+from onshape_robotics_toolkit.formats import MJCFSerializer
 from onshape_robotics_toolkit.graph import KinematicGraph
 from onshape_robotics_toolkit.models.document import Document
 from onshape_robotics_toolkit.parse import CAD
@@ -315,9 +316,10 @@ def find_best_design_variables(trial):
     graph = KinematicGraph.from_cad(cad, use_user_defined_root=True)
     ballbot = Robot.from_graph(kinematic_graph=graph, client=client, name="ballbot")
 
-    ballbot.set_robot_position(pos=(0, 0, 0.35))
-    ballbot = modify_ballbot(ballbot)
-    ballbot.save("ballbot.xml")
+    # Create config with simulation settings and save
+    config = create_ballbot_config(position=(0, 0, 0.35))
+    serializer = MJCFSerializer(config)
+    serializer.save(ballbot, "ballbot.xml")
 
     model = mujoco.MjModel.from_xml_path(filename="ballbot.xml")
     data = mujoco.MjData(model)

examples/simulation/mods.py

@@ -1,37 +1,50 @@
 from lxml import etree as ET
 
+from onshape_robotics_toolkit.formats import MJCFConfig, load_element
 from onshape_robotics_toolkit.models.mjcf import IMU, Gyro
-from onshape_robotics_toolkit.robot import Robot, load_element
 
 
-def modify_ballbot(ballbot: Robot) -> Robot:
-    ballbot.add_light(
+def create_ballbot_config(position: tuple[float, float, float] = (0, 0, 0.35)) -> MJCFConfig:
+    """
+    Create an MJCFConfig with all the ballbot simulation settings.
+
+    Args:
+        position: The robot position in world coordinates
+
+    Returns:
+        MJCFConfig configured for ballbot simulation
+    """
+    config = MJCFConfig(position=position)
+
+    # Add light
+    config.add_light(
         name="light_1",
         directional=True,
-        diffuse=[0.6, 0.6, 0.6],
-        specular=[0.2, 0.2, 0.2],
-        pos=[0, 0, 4],
-        direction=[0, 0, -1],
+        diffuse=(0.6, 0.6, 0.6),
+        specular=(0.2, 0.2, 0.2),
+        pos=(0, 0, 4),
+        direction=(0, 0, -1),
         castshadow=False,
     )
 
-    ballbot.add_actuator(
+    # Add actuators with encoders and force sensors
+    config.add_actuator(
         actuator_name="motor_1",
         joint_name="Revolute_1",
         ctrl_limited=True,
         ctrl_range=(-50, 50),
         add_encoder=True,
         add_force_sensor=True,
     )
-    ballbot.add_actuator(
+    config.add_actuator(
         actuator_name="motor_2",
         joint_name="Revolute_2",
         ctrl_limited=True,
         ctrl_range=(-50, 50),
         add_encoder=True,
         add_force_sensor=True,
     )
-    ballbot.add_actuator(
+    config.add_actuator(
         actuator_name="motor_3",
         joint_name="Revolute_3",
         ctrl_limited=True,
@@ -40,20 +53,21 @@ def modify_ballbot(ballbot: Robot) -> Robot:
         add_force_sensor=True,
     )
 
-    # For adding to specific named elements (like bodies)
+    # Add IMU site to specific body
     imu_site = ET.Element("site", name="imu", size="0.01", pos="0 0 0")
-    ballbot.add_custom_element_by_name("imu", "Part_3_1", imu_site)
+    config.add_custom_element_by_name("imu", "Part_3_1", imu_site)
 
-    # Add sensor
-    ballbot.add_sensor(
+    # Add sensors
+    config.add_sensor(
         name="imu",
         sensor=IMU(name="imu", objtype="site", objname="imu"),
     )
-    ballbot.add_sensor(
+    config.add_sensor(
         name="gyro_1",
         sensor=Gyro(name="gyro_1", site="imu"),
     )
 
+    # Create contact pairs
     contact = ET.Element("contact")
     pair_1 = ET.SubElement(contact, "pair")
     pair_1.set("geom1", "Part_2_3_collision")
@@ -90,20 +104,17 @@ def modify_ballbot(ballbot: Robot) -> Robot:
     pair_7.set("geom2", "floor")
     pair_7.set("friction", "1 10 3 10 10")
 
-    ballbot.add_custom_element_by_tag(name="contact", parent_tag="mujoco", element=contact)
+    config.add_custom_element_by_tag(name="contact", parent_tag="mujoco", element=contact)
 
+    # Add custom mesh and ball element
     ballbot_mesh = ET.Element("mesh", attrib={"name": "Part_1_1", "file": "meshes\\ball.stl"})
-    ballbot.add_custom_element_by_tag(name="ballbot", parent_tag="asset", element=ballbot_mesh)
+    config.add_custom_element_by_tag(name="ballbot", parent_tag="asset", element=ballbot_mesh)
     ball = load_element("ball.xml")
-    ballbot.add_custom_element_by_tag(name="ball", parent_tag="worldbody", element=ball)
-
-    # # set friction="1.0 0.01 0.001" for Part-2-1, Part-2-2, Part-2-3
-    # ballbot.set_element_attributes(element_name="Part-2-1-collision", attributes={"friction": "0.1 0.05 0.001"})
-    # ballbot.set_element_attributes(element_name="Part-2-2-collision", attributes={"friction": "0.1 0.05 0.001"})
-    # ballbot.set_element_attributes(element_name="Part-2-3-collision", attributes={"friction": "0.1 0.05 0.001"})
+    config.add_custom_element_by_tag(name="ball", parent_tag="worldbody", element=ball)
 
-    ballbot.set_element_attributes(element_name="Revolute_1", attributes={"axis": "0 0 1", "damping": "0.05"})
-    ballbot.set_element_attributes(element_name="Revolute_2", attributes={"axis": "0 0 1", "damping": "0.05"})
-    ballbot.set_element_attributes(element_name="Revolute_3", attributes={"axis": "0 0 1", "damping": "0.05"})
+    # Set joint attributes
+    config.set_element_attributes(element_name="Revolute_1", attributes={"axis": "0 0 1", "damping": "0.05"})
+    config.set_element_attributes(element_name="Revolute_2", attributes={"axis": "0 0 1", "damping": "0.05"})
+    config.set_element_attributes(element_name="Revolute_3", attributes={"axis": "0 0 1", "damping": "0.05"})
 
-    return ballbot
+    return config

onshape_robotics_toolkit/formats/__init__.py

@@ -27,7 +27,7 @@
 """
 
 from onshape_robotics_toolkit.formats.base import RobotDeserializer, RobotSerializer
-from onshape_robotics_toolkit.formats.mjcf import MJCFConfig, MJCFSerializer
+from onshape_robotics_toolkit.formats.mjcf import MJCFConfig, MJCFSerializer, load_element
 from onshape_robotics_toolkit.formats.urdf import URDFSerializer
 
 __all__ = [
@@ -36,4 +36,5 @@
     "RobotDeserializer",
     "RobotSerializer",
     "URDFSerializer",
+    "load_element",
 ]

onshape_robotics_toolkit/formats/mjcf.py

@@ -16,9 +16,28 @@
 
 from onshape_robotics_toolkit.formats.base import RobotDeserializer, RobotSerializer
 from onshape_robotics_toolkit.models.joint import BaseJoint
-from onshape_robotics_toolkit.models.mjcf import Actuator, Light, Sensor
+from onshape_robotics_toolkit.models.mjcf import Actuator, Encoder, ForceSensor, Light, Sensor
 from onshape_robotics_toolkit.utilities.helpers import format_number
 
+
+def load_element(file_name: str) -> ET._Element:
+    """
+    Load an XML element from a file.
+
+    Args:
+        file_name: The path to the XML file.
+
+    Returns:
+        The root element of the XML file.
+
+    Examples:
+        >>> element = load_element("ball.xml")
+        >>> config.add_custom_element_by_tag("ball", "worldbody", element)
+    """
+    tree: ET._ElementTree = ET.parse(file_name)  # noqa: S320
+    root: ET._Element = tree.getroot()
+    return root
+
 if TYPE_CHECKING:
     from onshape_robotics_toolkit.parse import PathKey
     from onshape_robotics_toolkit.robot import Robot
@@ -180,14 +199,17 @@ def set_element_attributes(
         self,
         element_name: str,
         attributes: dict[str, str],
-    ) -> None:
+    ) -> "MJCFConfig":
         """
         Register attribute modifications for an existing XML element.
 
         Args:
             element_name: The name of the element to modify
             attributes: Dictionary of attribute key-value pairs to set/update
 
+        Returns:
+            Self for method chaining
+
         Examples:
             >>> config = MJCFConfig()
             >>> config.set_element_attributes(
@@ -196,6 +218,128 @@ def set_element_attributes(
             ... )
         """
         self.mutated_elements[element_name] = attributes
+        return self
+
+    def add_light(
+        self,
+        name: str,
+        directional: bool,
+        diffuse: tuple[float, float, float],
+        specular: tuple[float, float, float],
+        pos: tuple[float, float, float],
+        direction: tuple[float, float, float],
+        castshadow: bool,
+    ) -> "MJCFConfig":
+        """
+        Add a light source to the scene.
+
+        Args:
+            name: The name of the light
+            directional: Whether the light is directional
+            diffuse: The diffuse color (r, g, b)
+            specular: The specular color (r, g, b)
+            pos: The position (x, y, z)
+            direction: The direction (x, y, z)
+            castshadow: Whether the light casts shadows
+
+        Returns:
+            Self for method chaining
+
+        Examples:
+            >>> config = MJCFConfig()
+            >>> config.add_light(
+            ...     name="sun",
+            ...     directional=True,
+            ...     diffuse=(0.8, 0.8, 0.8),
+            ...     specular=(0.2, 0.2, 0.2),
+            ...     pos=(0, 0, 10),
+            ...     direction=(0, 0, -1),
+            ...     castshadow=True
+            ... )
+        """
+        self.lights[name] = Light(
+            directional=directional,
+            diffuse=diffuse,
+            specular=specular,
+            pos=pos,
+            direction=direction,
+            castshadow=castshadow,
+        )
+        return self
+
+    def add_actuator(
+        self,
+        actuator_name: str,
+        joint_name: str,
+        ctrl_limited: bool = False,
+        ctrl_range: tuple[float, float] = (0, 0),
+        gear: float = 1.0,
+        add_encoder: bool = False,
+        add_force_sensor: bool = False,
+    ) -> "MJCFConfig":
+        """
+        Add an actuator to the model.
+
+        Args:
+            actuator_name: The name of the actuator
+            joint_name: The name of the joint to actuate
+            ctrl_limited: Whether the actuator has control limits
+            ctrl_range: The control range (min, max)
+            gear: The gear ratio
+            add_encoder: Whether to add an encoder sensor
+            add_force_sensor: Whether to add a force sensor
+
+        Returns:
+            Self for method chaining
+
+        Examples:
+            >>> config = MJCFConfig()
+            >>> config.add_actuator(
+            ...     actuator_name="motor1",
+            ...     joint_name="joint1",
+            ...     ctrl_limited=True,
+            ...     ctrl_range=(-10, 10),
+            ...     add_encoder=True,
+            ...     add_force_sensor=True
+            ... )
+        """
+        self.actuators[actuator_name] = Actuator(
+            name=actuator_name,
+            joint=joint_name,
+            ctrllimited=ctrl_limited,
+            ctrlrange=ctrl_range,
+            gear=gear,
+        )
+
+        if add_encoder:
+            encoder_name = f"{actuator_name}-enc"
+            self.sensors[encoder_name] = Encoder(encoder_name, actuator_name)
+
+        if add_force_sensor:
+            force_name = f"{actuator_name}-frc"
+            self.sensors[force_name] = ForceSensor(force_name, actuator_name)
+
+        return self
+
+    def add_sensor(self, name: str, sensor: Sensor) -> "MJCFConfig":
+        """
+        Add a sensor to the model.
+
+        Args:
+            name: The name of the sensor
+            sensor: The sensor object (IMU, Gyro, Encoder, ForceSensor)
+
+        Returns:
+            Self for method chaining
+
+        Examples:
+            >>> from onshape_robotics_toolkit.models.mjcf import IMU, Gyro
+            >>> config = MJCFConfig()
+            >>> config.add_sensor("imu", IMU(name="imu", objtype="site", objname="imu"))
+            >>> config.add_sensor("gyro", Gyro(name="gyro", site="imu"))
+        """
+        self.sensors[name] = sensor
+        return self
 
 
 class MJCFSerializer(RobotSerializer):

pyproject.toml

@@ -25,6 +25,12 @@ dependencies = [
 ]
 
 [project.optional-dependencies]
+simulation = [
+    "mujoco>=3.3.7",
+    "optuna>=4.6.0",
+    "plotly>=6.5.0",
+    "usd-core>=25.11",
+]
 docs = [
     "mkdocs>=1.4.2",
     "mkdocs-material>=9.2.7",