From 468c81aaa5f2f0c90ee273b2ebc08798a73decfe Mon Sep 17 00:00:00 2001 From: budzianowski Date: Fri, 7 Jun 2024 16:44:27 -0700 Subject: [PATCH] add reference position and limits --- sim/scripts/create_mjcf.py | 127 ++++++++++++++++++++++++------------- sim/stompy/joints.py | 125 ++++++++++++++++++++++++++++++++++++ 2 files changed, 207 insertions(+), 45 deletions(-) diff --git a/sim/scripts/create_mjcf.py b/sim/scripts/create_mjcf.py index cad140c2..0ae4c331 100644 --- a/sim/scripts/create_mjcf.py +++ b/sim/scripts/create_mjcf.py @@ -17,10 +17,9 @@ from pathlib import Path from typing import List, Union -from kol.formats import mjcf - from sim.env import model_dir, stompy_mjcf_path -from sim.stompy.joints import StompyFixed +from sim.scripts import mjcf +from sim.stompy.joints import MjcfStompy, StompyFixed logger = logging.getLogger(__name__) @@ -36,28 +35,36 @@ def _pretty_print_xml(xml_string: str) -> str: class Sim2SimRobot(mjcf.Robot): """A class to adapt the world in a Mujoco XML file.""" - def adapt_world(self) -> None: + def adapt_world( + self, add_floor: bool = True, add_reference_position: bool = False, remove_frc_range: bool = False + ) -> None: root: ET.Element = self.tree.getroot() - asset = root.find("asset") - asset.append( - ET.Element( - "texture", - name="texplane", - type="2d", - builtin="checker", - rgb1=".0 .0 .0", - rgb2=".8 .8 .8", - width="100", - height="108", + if add_floor: + asset = root.find("asset") + asset.append( + ET.Element( + "texture", + name="texplane", + type="2d", + builtin="checker", + rgb1=".0 .0 .0", + rgb2=".8 .8 .8", + width="100", + height="108", + ) ) - ) - asset.append( - ET.Element( - "material", name="matplane", reflectance="0.", texture="texplane", texrepeat="1 1", texuniform="true" + asset.append( + ET.Element( + "material", + name="matplane", + reflectance="0.", + texture="texplane", + texrepeat="1 1", + texuniform="true", + ) ) - ) - asset.append(ET.Element("material", name="visualgeom", rgba="0.5 0.9 0.2 1")) + asset.append(ET.Element("material", name="visualgeom", rgba="0.5 0.9 0.2 1")) compiler = root.find("compiler") if self.compiler is not None: @@ -104,20 +111,21 @@ def adapt_world(self) -> None: directional=True, diffuse=(0.6, 0.6, 0.6), specular=(0.2, 0.2, 0.2), pos=(0, 0, 4), dir=(0, 0, -1) ).to_xml(), ) - worldbody.insert( - 0, - mjcf.Geom( - name="ground", - type="plane", - size=(0, 0, 1), - pos=(0.001, 0, 0), - quat=(1, 0, 0, 0), - material="matplane", - condim=3, - conaffinity=1, - contype=0, - ).to_xml(), - ) + if add_floor: + worldbody.insert( + 0, + mjcf.Geom( + name="ground", + type="plane", + size=(0, 0, 1), + pos=(0.001, 0, 0), + quat=(1, 0, 0, 0), + material="matplane", + condim=3, + conaffinity=1, + contype=0, + ).to_xml(), + ) motors: List[mjcf.Motor] = [] sensor_pos: List[mjcf.Actuatorpos] = [] @@ -138,6 +146,8 @@ def adapt_world(self) -> None: sensor_vel.append(mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13")) sensor_frc.append(mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001)) + root = self.add_joint_limits(root, fixed=False) + # Add motors and sensors root.append(mjcf.Actuator(motors).to_xml()) root.append(mjcf.Sensor(sensor_pos, sensor_vel, sensor_frc).to_xml()) @@ -188,6 +198,9 @@ def adapt_world(self) -> None: ).to_xml(), ) + if add_reference_position: + root = self.add_reference_position(root) + # Move gathered elements to the new root body for item in items_to_move: worldbody.remove(item) @@ -216,12 +229,36 @@ def adapt_world(self) -> None: index = list(body).index(geom) body.insert(index + 1, new_geom) - # # Remove frc ranges? - # for body in root.findall(".//body"): - # joints = list(body.findall("joint")) - # for join in joints: - # if "actuatorfrcrange" in join.attrib: - # join.attrib.pop("actuatorfrcrange") + if remove_frc_range: + for body in root.findall(".//body"): + joints = list(body.findall("joint")) + for join in joints: + if "actuatorfrcrange" in join.attrib: + join.attrib.pop("actuatorfrcrange") + + def add_reference_position(self, root: ET.Element) -> None: + # Find all 'joint' elements + joints = root.findall(".//joint") + + default_standing = MjcfStompy.default_standing() + for joint in joints: + if joint.get("name") in default_standing.keys(): + joint.set("ref", str(default_standing[joint.get("name")])) + + return root + + def add_joint_limits(self, root: ET.Element, fixed: bool = False) -> None: + joint_limits = MjcfStompy.default_limits() + + for joint in root.findall(".//joint"): + joint_name = joint.get("name") + if joint_name in joint_limits: + limits = joint_limits.get(joint_name) + lower = str(limits.get("lower", 0.0)) + upper = str(limits.get("upper", 0.0)) + joint.set("range", f"{lower} {upper}") + + return root def save(self, path: Union[str, Path]) -> None: rough_string = ET.tostring(self.tree.getroot(), "utf-8") @@ -233,12 +270,12 @@ def save(self, path: Union[str, Path]) -> None: if __name__ == "__main__": - robot_name = "robot_fixed" + robot_name = "robot" robot = Sim2SimRobot( robot_name, model_dir(), mjcf.Compiler(angle="radian", meshdir="meshes", autolimits=True), - remove_inertia=False, + remove_inertia=True, ) - robot.adapt_world() - robot.save(stompy_mjcf_path(legs_only=True)) + robot.adapt_world(add_reference_position=True) + robot.save(stompy_mjcf_path(legs_only=False)) diff --git a/sim/stompy/joints.py b/sim/stompy/joints.py index 6f4a45a6..a87e9be7 100755 --- a/sim/stompy/joints.py +++ b/sim/stompy/joints.py @@ -292,6 +292,131 @@ def default_limits(cls) -> Dict[str, Dict[str, float]]: Stompy.legs.right.foot_roll: {"lower": -0.3, "upper": 0.3}, Stompy.legs.left.foot_roll: {"lower": -0.3, "upper": 0.3}, } + + +class ImuTorso(Torso): + pitch = "torso_1_x8_1_dof_x8" + pelvis_tx = "pelvis_tx" + pelvis_tz = "pelvis_tz" + pelvis_ty = "pelvis_ty" + tilt = "pelvis_tilt" + list = "pelvis_list" + rotation = "pelvis_rotation" + + +class MjcfStompy(Stompy): + head = Head() + torso = ImuTorso() + left_arm = LeftArm() + right_arm = RightArm() + legs = Legs() + + # test the model + @classmethod + def default_standing(cls) -> Dict[str, float]: + return { + MjcfStompy.head.left_right: 0.8526, + # arms + MjcfStompy.left_arm.shoulder_yaw: 1.57, + MjcfStompy.left_arm.shoulder_pitch: 1.59, + MjcfStompy.left_arm.elbow_roll: 0.0, + MjcfStompy.left_arm.elbow_yaw: 0.0, + + MjcfStompy.right_arm.elbow_roll: 2.48, + MjcfStompy.right_arm.elbow_yaw: -0.09, + + MjcfStompy.right_arm.shoulder_yaw: -1.6, + MjcfStompy.right_arm.shoulder_pitch: -1.64, + + # left hand + MjcfStompy.left_arm.hand.hand_roll: 0.9, + + # right hand + MjcfStompy.right_arm.hand.hand_roll: 0.64, + + # left leg + MjcfStompy.legs.left.hip_roll: -0.52, + MjcfStompy.legs.left.hip_yaw: 0.6, + MjcfStompy.legs.left.hip_pitch: -0.8, + MjcfStompy.legs.right.hip_roll: 0.59, + MjcfStompy.legs.right.hip_yaw: 0.7, + MjcfStompy.legs.right.hip_pitch: 1.0, + + # right leg + MjcfStompy.legs.left.knee: 0.2, + MjcfStompy.legs.right.knee: 0.0, + MjcfStompy.legs.left.ankle: 0.0, + MjcfStompy.legs.right.ankle: 0.0, + MjcfStompy.legs.left.foot_roll: 0.0, + MjcfStompy.legs.right.foot_roll: 0.0, + } + + @classmethod + def default_limits(cls) -> Dict[str, Dict[str, float]]: + return { + MjcfStompy.head.left_right: { + "lower": -3.14, + "upper": 3.14, + }, + MjcfStompy.right_arm.shoulder_yaw: { + "lower": -2.09, + "upper": 2.09, + }, + MjcfStompy.left_arm.shoulder_yaw: { + "lower": -2.09, + "upper": 2.09, + }, + MjcfStompy.right_arm.shoulder_pitch: { + "lower": -1.91, + "upper": 1.91, + }, + MjcfStompy.left_arm.shoulder_pitch: { + "lower": -1.91, + "upper": 1.91, + }, + MjcfStompy.legs.right.hip_roll: { + "lower": -1.04, + "upper": 1.04, + }, + MjcfStompy.legs.left.hip_roll: { + "lower": -1.04, + "upper": 1.04, + }, + MjcfStompy.legs.right.hip_yaw: { + "lower": -1.57, + "upper": 1.06, + }, + MjcfStompy.legs.left.hip_yaw: { + "lower": -1.57, + "upper": 1.06, + }, + MjcfStompy.legs.right.hip_pitch: { + "lower": -1.78, + "upper": 2.35, + }, + MjcfStompy.legs.left.hip_pitch: { + "lower": -1.78, + "upper": 2.35, + }, + MjcfStompy.legs.right.knee: { + "lower": 0, + "upper": 1.2, + }, + MjcfStompy.legs.left.knee: { + "lower": -1.2, + "upper": 0, + }, + MjcfStompy.legs.right.ankle: { + "lower": -0.3, + "upper": 0.3, + }, + MjcfStompy.legs.left.ankle: { + "lower": -0.3, + "upper": 0.3, + }, + MjcfStompy.legs.right.foot_roll: {"lower": -0.3, "upper": 0.3}, + MjcfStompy.legs.left.foot_roll: {"lower": -0.3, "upper": 0.3}, + } def print_joints() -> None: