Skip to content

Commit

Permalink
Adapt mjcf to the real environment (#11)
Browse files Browse the repository at this point in the history 
* save partial updates

* updates to the xml config, stompy is weirdly fixed

* all things in place, inertia has to be computed

* model behaves properly now

* add more options
  • Loading branch information
Paweł Budzianowski authored May 15, 2024
1 parent 63c5dad commit b252b4c
Show file tree
Hide file tree
Showing 2 changed files with 188 additions and 76 deletions.
263 changes: 187 additions & 76 deletions sim/scripts/create_mjcf.py
View file
Original file line number Diff line number Diff line change
@@ -1,78 +1,189 @@
# mypy: disable-error-code="valid-newtype"
"""This script updates the MJCF version for proper rendering."""
import os
# mypy: disable-error-code="operator,union-attr"
"""Defines common types and functions for exporting MJCF files.
API reference:
https://github.com/google-deepmind/mujoco/blob/main/src/xml/xml_native_writer.cc#L780
python sim.scripts/create_mjcf.py
Todo:
1. Add to all geoms
3. Condim 3 and 4 and difference in results
"""

import xml.etree.ElementTree as ET
from pathlib import Path
import xml.dom.minidom

from kol.formats import mjcf
from sim.stompy.joints import StompyFixed


STOMPY_HEIGHT = 1.2


def _pretty_print_xml(xml_string):
"""Formats the provided XML string into a pretty-printed version."""
parsed_xml = xml.dom.minidom.parseString(xml_string)
return parsed_xml.toprettyxml(indent=" ")


class Sim2SimRobot(mjcf.Robot):
"""A class to adapt the world in a Mujoco XML file."""

def adapt_world(self) -> None:
root = self.tree.getroot()

root.append(
mjcf.Option(
timestep=0.001,
viscosity=1e-6,
iterations=50,
solver="PGS",
gravity=(0, 0, -9.81),
flag=mjcf.Flag(frictionloss="enable"),
).to_xml()
)

# TODO - test the physical parameters
root.append(
mjcf.Default(
joint=mjcf.Joint(armature=0.01, damping=0.1, limited=True, frictionloss=0.01),
motor=mjcf.Motor(ctrllimited=True),
equality=mjcf.Equality(solref=(0.001, 2)),
geom=mjcf.Geom(
solref=(0.001, 2),
friction=(0.9, 0.2, 0.2),
condim=4,
contype=1,
conaffinity=15,
),
# visualgem?
# joint param damping
).to_xml()
)
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"
)
)

compiler = root.find("compiler")
if self.compiler is not None:
compiler = self.compiler.to_xml(compiler)

worldbody = root.find("worldbody")
# List to store items to be moved to the new root body
items_to_move = []
# Gather all children (geoms and bodies) that need to be moved under the new root body
for element in worldbody:
items_to_move.append(element)

new_root_body = mjcf.Body(name="root", pos=(0, 0, STOMPY_HEIGHT), quat=(1, 0, 0, 0)).to_xml()

# Add joints to all the movement of the base
new_root_body.extend(
[
mjcf.Joint(name="root_x", type="slide", axis=(1, 0, 0), limited=False).to_xml(),
mjcf.Joint(name="root_y", type="slide", axis=(0, 1, 0), limited=False).to_xml(),
mjcf.Joint(name="root_z", type="slide", axis=(0, 0, 1), limited=False).to_xml(),
mjcf.Joint(name="root_ball", type="ball", limited=False).to_xml(),
]
)

# Add imu site to the body - relative position to the body
# check at what stage we use this
new_root_body.append(mjcf.Site(name="imu", size=0.01, pos=(0, 0, 0)).to_xml())

# Move gathered elements to the new root body
for item in items_to_move:
worldbody.remove(item)
new_root_body.append(item)

# Add the new root body to the worldbody
worldbody.append(new_root_body)

worldbody.insert(
0,
mjcf.Light(
directional=True,
diffuse=(0.4, 0.4, 0.4),
specular=(0.1, 0.1, 0.1),
pos=(0, 0, 5.0),
dir=(0, 0, -1),
castshadow=False,
).to_xml(),
)
worldbody.insert(
0,
mjcf.Light(
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=1,
conaffinity=15,
).to_xml(),
)

motors = []
sensors = []
for joint, limits in StompyFixed.default_limits().items():
if joint in StompyFixed.legs.all_joints():
motors.append(
mjcf.Motor(
name=joint, joint=joint, gear=1, ctrlrange=(limits["lower"], limits["upper"]), ctrllimited=True
)
)
sensors.extend(
[
mjcf.Actuatorpos(name=joint + "_p", actuator=joint, user="13"),
mjcf.Actuatorvel(name=joint + "_v", actuator=joint, user="13"),
mjcf.Actuatorfrc(name=joint + "_f", actuator=joint, user="13", noise=0.001),
]
)

# Add motors and sensors
root.append(mjcf.Actuator(motors).to_xml())
root.append(mjcf.Sensor(sensors).to_xml())

# Add imus
sensors = root.find("sensor")
sensors.extend(
[
# TODO - pfb30 test that
# ET.Element("framequat", name="orientation", objtype="site", noise="0.001", objname="imu"),
# ET.Element("framepos", name="position", objtype="site", noise="0.001", objname="imu"),
ET.Element("gyro", name="angular-velocity", site="imu", noise="0.005", cutoff="34.9"),
ET.Element("velocimeter", name="linear-velocity", site="imu", noise="0.001", cutoff="30"),
ET.Element("accelerometer", name="linear-acceleration", site="imu", noise="0.005", cutoff="157"),
ET.Element("magnetometer", name="magnetometer", site="imu"),
]
)
self.tree = ET.ElementTree(root)

import mujoco


def create_mjcf() -> None:
model_xml = mujoco.MjModel.from_xml_path(os.path.join(os.getenv("MODEL_DIR"), "robot.urdf"))
mujoco.mj_saveLastXML(os.path.join(os.getenv("MODEL_DIR"), "robot.xml"), model_xml)

# Add light and floor to the xml file and move all elements to a new root body
tree = ET.parse(os.path.join(os.getenv("MODEL_DIR"), "robot.xml"))
root = tree.getroot()
# Create light element
light = ET.Element(
"light",
{
"cutoff": "100",
"diffuse": "1 1 1",
"dir": "-0 0 -1.3",
"directional": "true",
"exponent": "1",
"pos": "0 0 1.3",
"specular": ".1 .1 .1",
},
)

# Create floor geometry element
floor = ET.Element(
"geom",
{
"conaffinity": "1",
"condim": "3",
"name": "floor",
"pos": "0 0 -1.5",
"rgba": "0.8 0.9 0.8 1",
"size": "40 40 40",
"type": "plane",
# "material": "MatPlane",
},
)
# Access the <worldbody> element
worldbody = root.find("worldbody")
# Insert new elements at the beginning of the worldbody
worldbody.insert(0, floor)
worldbody.insert(0, light)

new_root_body = ET.Element("body", name="root", pos="-1.4 0 -0.3", quat="1 0 0 1")

# List to store items to be moved to the new root body
items_to_move = []
# Gather all children (geoms and bodies) that need to be moved under the new root body
for element in worldbody:
items_to_move.append(element)
# Move gathered elements to the new root body
for item in items_to_move:
worldbody.remove(item)
new_root_body.append(item)
# Add the new root body to the worldbody
worldbody.append(new_root_body)
worldbody2 = worldbody
# # Create a new root worldbody and add the new root body to it
# new_worldbody = ET.Element("worldbody", name="new_root")
index = list(root).index(worldbody)

# # Remove the old <worldbody> and insert the new one at the same index
root.remove(worldbody)
root.insert(index, worldbody2)
modified_xml = ET.tostring(root, encoding="unicode")

with open(os.path.join(os.getenv("MODEL_DIR"), "robot.xml"), "w") as f:
f.write(modified_xml)


if __name__ == "__main__":
create_mjcf()
def save(self, path: str | Path) -> None:
rough_string = ET.tostring(self.tree.getroot(), "utf-8")
# Pretty print the XML
formatted_xml = _pretty_print_xml(rough_string)
1 change: 1 addition & 0 deletions sim/stompy/joints.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,7 @@ def default_standing(cls) -> Dict[str, float]:
Stompy.legs.right.foot_roll: np.deg2rad(0), # 0
}

@classmethod
def default_limits(cls) -> Dict[str, Dict[str, float]]:
return {
Stompy.head.left_right: {
Expand Down

0 comments on commit b252b4c

Please sign in to comment.