Skip to content

Commit

Permalink
add reference position and limits (#16)
Browse files Browse the repository at this point in the history
  • Loading branch information
Paweł Budzianowski authored Jun 8, 2024
1 parent feded53 commit 5b910ba
Show file tree
Hide file tree
Showing 2 changed files with 207 additions and 45 deletions.
127 changes: 82 additions & 45 deletions sim/scripts/create_mjcf.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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__)

Expand All @@ -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:
Expand Down Expand Up @@ -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] = []
Expand All @@ -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())
Expand Down Expand Up @@ -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)
Expand Down Expand Up @@ -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")
Expand All @@ -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))
125 changes: 125 additions & 0 deletions sim/stompy/joints.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -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:
Expand Down

0 comments on commit 5b910ba

Please sign in to comment.