Added modifications to ballbot sim example.

neurobionics · Jan 13, 2025 · 9835bef · 9835bef
1 parent c8f47aa
commit 9835bef
Show file tree

Hide file tree

Showing 8 changed files with 139 additions and 109 deletions.
diff --git a/examples/simulation/ball.xml b/examples/simulation/ball.xml
@@ -1,6 +1,6 @@
 <body name="ball" pos="0 0 0.132" quat="1 0 0 0">
-    <freejoint name="ball" />
-    <inertial pos="0 0 -0.0075" quat="0.5 0.5 -0.5 0.5" mass="0.524" diaginertia="0.073541512 0.07356916 0.073543931" />
-    <geom name="Part-1-1-collision" type="mesh" rgba="1 0.55 0 1" mesh="Part-1-1" condim="3" contype="1" conaffinity="1" density="0" group="0" friction="0.1 0.05 0.001"/>
+    <freejoint name="ball"/>
+    <inertial pos="0 0 -0.0075" quat="0.5 0.5 -0.5 0.5" mass="0.524" diaginertia="0.073541512 0.07356916 0.073543931"/>
+    <geom name="Part-1-1-collision" type="sphere" size="0.132" rgba="1 0.55 0 1" contype="1" conaffinity="1" density="0" group="0"/>
     <geom type="mesh" rgba="1 0.55 0 1" mesh="Part-1-1" conaffinity="0" condim="1" contype="0" density="0" group="1"/>
-</body>
+</body>
diff --git a/examples/simulation/main.py b/examples/simulation/main.py
@@ -16,81 +16,80 @@
 WIDTH = 640
 
 FREQUENCY = 200
+dt = 1 / FREQUENCY
 PHASE = 3
 
 # Variable bounds (in mm and degrees)
-WHEEL_DIAMETER_BOUNDS = (100, 200)
+WHEEL_DIAMETER_BOUNDS = (100, 120)
+SPACER_HEIGHT_BOUNDS = (65, 120)
 ALPHA_BOUNDS = (30, 55)
 
-
 SIMULATION_DURATION = 10  # seconds to run each trial
 MIN_HEIGHT = 0.15  # minimum height before considering failure
 
 # PID parameters for roll, pitch, and yaw
-KP = 2
-KI = 0.0
-KD = 0.0
-
-# Initialize integral and previous error for PID
-integral_roll = 0
-integral_pitch = 0
-
-prev_error_roll = 0
-prev_error_pitch = 0
-
-
-def get_theta(data):
-    rot = Rotation.from_quat(data.sensor("imu").data)
-    theta = rot.as_euler("xyz", degrees=False)
+KP = 160.0
+KI = 10.0
+KD = 2.0
+
+integral_error_x = 0.0
+integral_error_y = 0.0
+previous_error_x = 0.0
+previous_error_y = 0.0
+
+def get_theta(data, degrees=False):
+    rot = Rotation.from_quat(data.qpos[3:7], scalar_first=True)
+    theta = rot.as_euler("XYZ", degrees=degrees)
     return theta[0], theta[1], theta[2]
 
+def control(data, alpha):
+    global integral_error_x, integral_error_y, previous_error_x, previous_error_y
 
-def control(data, roll_sp=0, pitch_sp=0, yaw_sp=0):
-    global integral_roll, integral_pitch
-    global prev_error_roll, prev_error_pitch
-
-    roll, pitch, yaw = get_theta(data)
-    roll = roll - np.pi
-
-    # swap roll and pitch
-    roll, pitch = pitch, roll
+    # Get current orientation
+    theta = get_theta(data)
 
     # Calculate errors
-    error_roll = roll_sp - roll
-    error_pitch = pitch_sp - pitch
+    error_x = 0.0 - theta[0]
+    error_y = 0.0 - theta[1]
 
-    # Update integrals
-    integral_roll += error_roll
-    integral_pitch += error_pitch
+    # Update integral errors
+    integral_error_x += error_x * dt
+    integral_error_y += error_y * dt
 
-    # Calculate derivatives
-    derivative_roll = error_roll - prev_error_roll
-    derivative_pitch = error_pitch - prev_error_pitch
+    # Calculate derivative errors
+    derivative_error_x = (error_x - previous_error_x) / dt
+    derivative_error_y = (error_y - previous_error_y) / dt
 
-    # PID control for each axis
-    tx = KP * error_roll + KI * integral_roll + KD * derivative_roll
-    ty = -1 * (KP * error_pitch + KI * integral_pitch + KD * derivative_pitch)
+    # PID control calculations
+    tx = KP * error_x + KI * integral_error_x + KD * derivative_error_x
+    ty = KP * error_y + KI * integral_error_y + KD * derivative_error_y
     tz = 0.0
 
-    # Compute motor torques
-    t1, t2, t3 = compute_motor_torques(tx, ty, tz)
+    # Saturate the torque values to be within [-3, 3] Nm
+    tx = np.clip(tx, -3.0, 3.0)
+    ty = np.clip(ty, -3.0, 3.0)
+    tz = np.clip(tz, -3.0, 3.0)
 
-    # Apply control signals
+    # Update previous errors
+    previous_error_x = error_x
+    previous_error_y = error_y
+
+    # Compute motor torques
+    t1, t2, t3 = compute_motor_torques(np.deg2rad(alpha), tx, ty, tz)
     data.ctrl[0] = t1
     data.ctrl[1] = t2
     data.ctrl[2] = t3
 
-    # Update previous errors
-    prev_error_roll = error_roll
-    prev_error_pitch = error_pitch
-
 
 def objective(trial):
+    # reset global PID error values
     wheel_diameter = trial.suggest_float("wheel_diameter", WHEEL_DIAMETER_BOUNDS[0], WHEEL_DIAMETER_BOUNDS[1])
+    spacer_height = trial.suggest_float("spacer_height", SPACER_HEIGHT_BOUNDS[0], SPACER_HEIGHT_BOUNDS[1])
     alpha = trial.suggest_float("alpha", ALPHA_BOUNDS[0], ALPHA_BOUNDS[1])
 
     variables["wheel_diameter"].expression = f"{wheel_diameter:.1f} mm"
     variables["alpha"].expression = f"{alpha:.1f} deg"
+    variables["spacer_height"].expression = f"{spacer_height:.1f} mm"
     client.set_variables(doc.did, doc.wid, elements["variables"].id, variables)
 
     ballbot: Robot = Robot.from_url(
@@ -114,12 +113,21 @@ def objective(trial):
             timesteps = 0
             max_timesteps = int(SIMULATION_DURATION / model.opt.timestep)
             total_angle_error = 0
+            # reset global PID error values
+            global integral_error_x, integral_error_y, previous_error_x, previous_error_y
+            integral_error_x = 0.0
+            integral_error_y = 0.0
+            previous_error_x = 0.0
+            previous_error_y = 0.0
+
             # Reset data for a new trial
             mujoco.mj_resetData(model, data)
 
             while timesteps < max_timesteps and viewer.is_running():
                 mujoco.mj_step(model, data)
-                control(data)
+
+                if data.time > 0.3:
+                    control(data, alpha)
 
                 if data.body("ballbot").xpos[2] < MIN_HEIGHT:
                     print("Ballbot fell below minimum height, ending trial.")

diff --git a/examples/simulation/mods.py b/examples/simulation/mods.py
@@ -27,24 +27,24 @@ def modify_ballbot(ballbot: Robot) -> Robot:
     ballbot.add_actuator(
         actuator_name="motor-1",
         joint_name="Revolute-1",
-        ctrl_limited=False,
-        gear=70,
+        ctrl_limited=True,
+        ctrl_range=(-3, 3),
         add_encoder=True,
         add_force_sensor=True,
     )
     ballbot.add_actuator(
         actuator_name="motor-2",
         joint_name="Revolute-2",
-        ctrl_limited=False,
-        gear=70,
+        ctrl_limited=True,
+        ctrl_range=(-3, 3),
         add_encoder=True,
         add_force_sensor=True,
     )
     ballbot.add_actuator(
         actuator_name="motor-3",
         joint_name="Revolute-3",
-        ctrl_limited=False,
-        gear=70,
+        ctrl_limited=True,
+        ctrl_range=(-3, 3),
         add_encoder=True,
         add_force_sensor=True,
     )
@@ -63,6 +63,24 @@ def modify_ballbot(ballbot: Robot) -> Robot:
         sensor=Gyro(name="gyro-1", site="imu", noise=0.001, cutoff=34.9),
     )
 
+    contact = ET.Element("contact")
+    pair_1 = ET.SubElement(contact, "pair")
+    pair_1.set("geom1", "Part-2-3-collision")
+    pair_1.set("geom2", "Part-1-1-collision")
+    pair_1.set("friction", "0.3 0.3 0.005 0.9 0.9")
+
+    pair_2 = ET.SubElement(contact, "pair")
+    pair_2.set("geom1", "Part-2-2-collision")
+    pair_2.set("geom2", "Part-1-1-collision")
+    pair_2.set("friction", "0.3 0.3 0.005 0.9 0.9")
+
+    pair_3 = ET.SubElement(contact, "pair")
+    pair_3.set("geom1", "Part-2-1-collision")
+    pair_3.set("geom2", "Part-1-1-collision")
+    pair_3.set("friction", "0.3 0.3 0.005 0.9 0.9")
+
+    ballbot.add_custom_element_by_tag(name="contact", parent_tag="mujoco", element=contact)
+
     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)
     ball = load_element("ball.xml")
@@ -73,9 +91,8 @@ def modify_ballbot(ballbot: Robot) -> Robot:
     # 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"})
 
-    # set force range to -1 to 1 for all joints
-    ballbot.set_element_attributes(element_name="Revolute-1", attributes={"actuatorfrcrange": "-3 3"})
-    ballbot.set_element_attributes(element_name="Revolute-2", attributes={"actuatorfrcrange": "-3 3"})
-    ballbot.set_element_attributes(element_name="Revolute-3", attributes={"actuatorfrcrange": "-3 3"})
+    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"})
 
     return ballbot
diff --git a/examples/simulation/transformations.py b/examples/simulation/transformations.py
@@ -1,3 +1,6 @@
+import numpy as np
+
+
 def transform_w2b(m1, m2, m3):
     """
     Returns Ball (Phi) attributes
@@ -10,10 +13,11 @@ def transform_w2b(m1, m2, m3):
     return x, y, z
 
 
-def compute_motor_torques(Tx, Ty, Tz):
+def compute_motor_torques(alpha, Tx, Ty, Tz):
     """
     Parameters:
     ----------
+    alpha: angle of the ball
     Tx: Torque along x-axis
     Ty: Torque along y-axis
     Tz: Torque along z-axis
@@ -25,7 +29,7 @@ def compute_motor_torques(Tx, Ty, Tz):
             |
             |
             |
-            . _ _ _ _ Ty
+    Ty_ _ _ .
            / \
           /   \
          /     \
@@ -37,8 +41,8 @@ def compute_motor_torques(Tx, Ty, Tz):
     T3: Motor Torque 3
     """
 
-    T1 = (0.3333) * (Tz - (2.8284 * Tx))
-    T2 = (0.3333) * (Tz + (1.4142 * (Tx - 1.7320 * Ty)))
-    T3 = (0.3333) * (Tz + (1.4142 * (Tx + 1.7320 * Ty)))
+    T1 = (0.3333) * (Tz + ((2/np.cos(alpha)) * Tx))
+    T2 = (0.3333) * (Tz + ((1/np.cos(alpha)) * (1.7320 * Ty - Tx)))
+    T3 = (0.3333) * (Tz - ((1/np.cos(alpha)) * (Tx + 1.7320 * Ty)))
 
     return T1, T2, T3
diff --git a/onshape_robotics_toolkit/models/joint.py b/onshape_robotics_toolkit/models/joint.py
@@ -417,8 +417,8 @@ class RevoluteJoint(BaseJoint):
         <Element 'joint' at 0x7f8b3c0b4c70>
     """
 
-    limits: JointLimits
     axis: Axis
+    limits: JointLimits | None = None
     dynamics: JointDynamics | None = None
     mimic: JointMimic | None = None
 
@@ -474,7 +474,8 @@ def to_mjcf(self, root: ET.Element) -> None:
         joint.set("pos", " ".join(map(str, self.origin.xyz)))
 
         self.axis.to_mjcf(joint)
-        joint.set("range", " ".join(map(str, [self.limits.lower, self.limits.upper])))
+        if self.limits:
+            joint.set("range", " ".join(map(str, [self.limits.lower, self.limits.upper])))
 
         if self.dynamics:
             joint.set("damping", str(self.dynamics.damping))

diff --git a/onshape_robotics_toolkit/models/mjcf.py b/onshape_robotics_toolkit/models/mjcf.py
@@ -106,8 +106,8 @@ class Actuator:
 
     name: str
     joint: str
-    ctrllimited: bool
-    gear: float
+    ctrllimited: bool = False
+    gear: float = 1.0
     ctrlrange: tuple[float, float] = (0, 0)
 
     def to_mjcf(self, root: ET.Element) -> None:

diff --git a/onshape_robotics_toolkit/robot.py b/onshape_robotics_toolkit/robot.py
@@ -58,7 +58,7 @@
     # "meshdir": "meshes",
 }
 
-DEFAULT_OPTION_ATTRIBUTES = {"timestep": "0.001", "gravity": "0 0 -9.81", "iterations": "50", "solver": "PGS"}
+DEFAULT_OPTION_ATTRIBUTES = {"timestep": "0.001", "gravity": "0 0 -9.81", "iterations": "50"}
 
 URDF_EULER_SEQ = "xyz"  # URDF uses XYZ fixed angles
 MJCF_EULER_SEQ = "XYZ"  # MuJoCo uses XYZ extrinsic rotations, capitalization matters
@@ -250,11 +250,11 @@ def add_actuator(
         self,
         actuator_name: str,
         joint_name: str,
-        ctrl_limited: bool,
-        gear: float,
+        ctrl_limited: bool = False,
         add_encoder: bool = True,
         add_force_sensor: bool = True,
         ctrl_range: tuple[float, float] = (0, 0),
+        gear: float = 1.0,
     ) -> None:
         """
         Add an actuator to the robot model.
@@ -664,7 +664,7 @@ def to_mjcf(self) -> str:  # noqa: C901
                 element = element_info["element"]
 
                 if find_by_tag:
-                    parent_element = model.find(parent)
+                    parent_element = model if parent == "mujoco" else model.find(parent)
                 else:
                     xpath = f".//body[@name='{parent}']"
                     parent_element = model.find(xpath)