Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

add reference position and limits #16

Merged
merged 1 commit into from
Jun 8, 2024
Merged

add reference position and limits #16

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
127 changes: 82 additions & 45 deletions sim/scripts/create_mjcf.py
View file
Open in desktop
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
Open in desktop
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
Loading