Buckets:

mishig
/

so101-block-sorting

Files

xet

mishig/so101-block-sorting / scripts /build_lego_scene.py

mishig

about 1 month ago

download

raw

5.92 kB

"""Create a physical block-sorting scene around the official SO-101 model."""
from pathlib import Path
import xml.etree.ElementTree as ET
import numpy as np

ROOT = Path(__file__).resolve().parent.parent
OUT = ROOT / "scene"
COLORS = {"red": ".88 .12 .10 1", "green": ".12 .70 .26 1", "blue": ".10 .32 .92 1"}
FIDUCIALS = [(0.11, -0.16), (0.37, -0.16), (0.37, 0.23), (0.11, 0.23)]


def main():
    source = ROOT / "assets/so101"
    parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True))
    robot = ET.parse(source / "so101_new_calib.xml", parser=parser)
    robot.getroot().find("compiler").set("meshdir", "../assets/so101/assets")
    for geom in robot.findall(".//geom[@class='collision']"):
        geom.set("friction", "1.3 .02 .002")
        geom.set("condim", "4")
    grip = robot.find(".//body[@name='gripper']")
    jaw = robot.find(".//body[@name='moving_jaw_so101_v1']")
    # The upstream convex hulls bridge the open jaw interior. Replace only
    # the fingertip collision surfaces with visible compliant rubber pads.
    for body, mesh in ((grip, "wrist_roll_follower_so101_v1"), (jaw, "moving_jaw_so101_v1")):
        for geom in list(body.findall("geom")):
            if geom.get("class") == "collision" and geom.get("mesh") == mesh:
                body.remove(geom)
    for body, name, pos, size in (
        (grip, "fixed_finger_pad", "-.0125 -.0002 -.091", ".0046 .008 .0135"),
        (jaw, "moving_finger_pad", "-.0077 -.0676 .019", ".0046 .0135 .008"),
    ):
        ET.SubElement(body, "geom", name=name, type="box", pos=pos, size=size,
                      rgba=".10 .11 .12 1", friction="1.5 .03 .003", condim="4", solref=".004 1", mass="0")
    ET.SubElement(grip, "site", name="grasp_center", pos=".0041 -.0002 -.093", size=".003", rgba="1 0 1 1", group="4")
    ET.SubElement(grip, "geom", name="vision_marker", type="sphere", pos=".0041 -.0002 -.078", size=".004", rgba="1 0 1 1", contype="0", conaffinity="0", mass="0")
    actuator = robot.find(".//actuator/position[@name='gripper']")
    if actuator is not None:
        actuator.set("forcerange", "-.4 .4")
    OUT.mkdir(parents=True, exist_ok=True)
    header = """<!-- Modified for this SO-101 block-sorting package: portable mesh paths,
     fingertip collision pads, grasp and vision markers, contact friction,
     and gripper force limit.
     Derived from TheRobotStudio/SO-ARM100, Simulation/SO101/so101_new_calib.xml,
     commit eecbe3e0a9ebb23e25ad7b2759b03884c6660903. Apache-2.0.
     Upstream generated using onshape-to-robot.
     Onshape https://cad.onshape.com/documents/7715cc284bb430fe6dab4ffd/w/4fd0791b683777b02f8d975a/e/826c553ede3b7592eb9ca800
     See THIRD_PARTY_NOTICES.md for license and source attribution. -->
"""
    (OUT / "so101.xml").write_text(header + ET.tostring(robot.getroot(), encoding="unicode"))
    scene = ET.Element("mujoco", model="SO-101 block sorting")
    ET.SubElement(scene, "include", file="so101.xml")
    ET.SubElement(scene, "option", timestep=".002", integrator="implicitfast", iterations="80", gravity="0 0 -9.81", noslip_iterations="10", impratio="10")
    visual = ET.SubElement(scene, "visual")
    ET.SubElement(visual, "global", offwidth="960", offheight="720")
    ET.SubElement(visual, "headlight", diffuse=".7 .7 .7", ambient=".5 .5 .5", specular=".1 .1 .1")
    asset = ET.SubElement(scene, "asset")
    ET.SubElement(asset, "texture", type="skybox", builtin="gradient", rgb1=".18 .23 .30", rgb2=".42 .48 .54", width="512", height="3072")
    world = ET.SubElement(scene, "worldbody")
    ET.SubElement(world, "light", pos=".1 -.3 1", dir="0 0 -1", directional="true", diffuse=".8 .8 .8")
    ET.SubElement(world, "geom", name="floor", type="plane", size="1 1 .02", rgba=".18 .22 .28 1", friction="1 .01 .001")
    ET.SubElement(world, "geom", name="work_mat", type="box", pos=".235 .035 .0005", size=".155 .215 .0005", rgba=".68 .70 .67 1")
    for i, (x, y) in enumerate(FIDUCIALS):
        ET.SubElement(world, "geom", name=f"calibration_{i}", type="cylinder", pos=f"{x} {y} .0012", size=".007 .0003", rgba=".05 .95 .95 1", contype="0", conaffinity="0")
    camera_pos = np.array([.70, -.65, .68])
    target = np.array([.18, .02, .045])
    z = camera_pos - target
    z /= np.linalg.norm(z)
    x = np.cross([0., 0., 1.], z); x /= np.linalg.norm(x)
    y = np.cross(z, x)
    ET.SubElement(world, "camera", name="front", pos=" ".join(map(str, camera_pos)), xyaxes=" ".join(map(str, np.r_[x, y])), fovy="42")
    placements = [("red", .255, -.105), ("green", .285, -.025), ("blue", .19, -.075)]
    for name, x, y in placements:
        body = ET.SubElement(world, "body", name=f"block_{name}", pos=f"{x} {y} .014")
        ET.SubElement(body, "freejoint", name=f"block_{name}_free")
        ET.SubElement(body, "geom", name=f"block_{name}_solid", type="box", size=".012 .010 .012", rgba=COLORS[name], mass=".014", friction="1.5 .03 .003", condim="4", solref=".008 1")
        for sx in (-.006, .006):
            for sy in (-.005, .005):
                ET.SubElement(body, "geom", type="cylinder", pos=f"{sx} {sy} .014", size=".003 .002", rgba=COLORS[name], mass=".0001", friction="1.2 .02 .002")
    box = ET.SubElement(world, "body", name="box", pos=".20 .135 0")
    ET.SubElement(box, "geom", name="box_bottom", type="box", pos="0 0 .003", size=".062 .050 .003", rgba=".52 .36 .19 1", friction="1 .02 .002")
    for name, pos, size in [
        ("left", "-.064 0 .020", ".002 .054 .020"),
        ("right", ".064 0 .020", ".002 .054 .020"),
        ("back", "0 .052 .020", ".062 .002 .020"),
        ("front", "0 -.052 .020", ".062 .002 .020"),
    ]:
        ET.SubElement(box, "geom", name=f"box_{name}", type="box", pos=pos, size=size, rgba=".66 .45 .24 1", friction="1 .02 .002")
    ET.indent(scene)
    ET.ElementTree(scene).write(OUT / "scene.xml", encoding="unicode")
    print(OUT / "scene.xml")


if __name__ == "__main__":
    main()

Xet Storage Details

Size:: 5.92 kB
Xet hash:: 39efd27c97dae026b1937abc8f466fe31863b4fa7e69b6a91dd592d612bc3b71

Xet efficiently stores files, intelligently splitting them into unique chunks and accelerating uploads and downloads. More info.

	"""Create a physical block-sorting scene around the official SO-101 model."""
	from pathlib import Path
	import xml.etree.ElementTree as ET
	import numpy as np

	ROOT = Path(__file__).resolve().parent.parent
	OUT = ROOT / "scene"
	COLORS = {"red": ".88 .12 .10 1", "green": ".12 .70 .26 1", "blue": ".10 .32 .92 1"}
	FIDUCIALS = [(0.11, -0.16), (0.37, -0.16), (0.37, 0.23), (0.11, 0.23)]


	def main():
	source = ROOT / "assets/so101"
	parser = ET.XMLParser(target=ET.TreeBuilder(insert_comments=True))
	robot = ET.parse(source / "so101_new_calib.xml", parser=parser)
	robot.getroot().find("compiler").set("meshdir", "../assets/so101/assets")
	for geom in robot.findall(".//geom[@class='collision']"):
	geom.set("friction", "1.3 .02 .002")
	geom.set("condim", "4")
	grip = robot.find(".//body[@name='gripper']")
	jaw = robot.find(".//body[@name='moving_jaw_so101_v1']")
	# The upstream convex hulls bridge the open jaw interior. Replace only
	# the fingertip collision surfaces with visible compliant rubber pads.
	for body, mesh in ((grip, "wrist_roll_follower_so101_v1"), (jaw, "moving_jaw_so101_v1")):
	for geom in list(body.findall("geom")):
	if geom.get("class") == "collision" and geom.get("mesh") == mesh:
	body.remove(geom)
	for body, name, pos, size in (
	(grip, "fixed_finger_pad", "-.0125 -.0002 -.091", ".0046 .008 .0135"),
	(jaw, "moving_finger_pad", "-.0077 -.0676 .019", ".0046 .0135 .008"),
	):
	ET.SubElement(body, "geom", name=name, type="box", pos=pos, size=size,
	rgba=".10 .11 .12 1", friction="1.5 .03 .003", condim="4", solref=".004 1", mass="0")
	ET.SubElement(grip, "site", name="grasp_center", pos=".0041 -.0002 -.093", size=".003", rgba="1 0 1 1", group="4")
	ET.SubElement(grip, "geom", name="vision_marker", type="sphere", pos=".0041 -.0002 -.078", size=".004", rgba="1 0 1 1", contype="0", conaffinity="0", mass="0")
	actuator = robot.find(".//actuator/position[@name='gripper']")
	if actuator is not None:
	actuator.set("forcerange", "-.4 .4")
	OUT.mkdir(parents=True, exist_ok=True)
	header = """<!-- Modified for this SO-101 block-sorting package: portable mesh paths,
	fingertip collision pads, grasp and vision markers, contact friction,
	and gripper force limit.
	Derived from TheRobotStudio/SO-ARM100, Simulation/SO101/so101_new_calib.xml,
	commit eecbe3e0a9ebb23e25ad7b2759b03884c6660903. Apache-2.0.
	Upstream generated using onshape-to-robot.
	Onshape https://cad.onshape.com/documents/7715cc284bb430fe6dab4ffd/w/4fd0791b683777b02f8d975a/e/826c553ede3b7592eb9ca800
	See THIRD_PARTY_NOTICES.md for license and source attribution. -->
	"""
	(OUT / "so101.xml").write_text(header + ET.tostring(robot.getroot(), encoding="unicode"))
	scene = ET.Element("mujoco", model="SO-101 block sorting")
	ET.SubElement(scene, "include", file="so101.xml")
	ET.SubElement(scene, "option", timestep=".002", integrator="implicitfast", iterations="80", gravity="0 0 -9.81", noslip_iterations="10", impratio="10")
	visual = ET.SubElement(scene, "visual")
	ET.SubElement(visual, "global", offwidth="960", offheight="720")
	ET.SubElement(visual, "headlight", diffuse=".7 .7 .7", ambient=".5 .5 .5", specular=".1 .1 .1")
	asset = ET.SubElement(scene, "asset")
	ET.SubElement(asset, "texture", type="skybox", builtin="gradient", rgb1=".18 .23 .30", rgb2=".42 .48 .54", width="512", height="3072")
	world = ET.SubElement(scene, "worldbody")
	ET.SubElement(world, "light", pos=".1 -.3 1", dir="0 0 -1", directional="true", diffuse=".8 .8 .8")
	ET.SubElement(world, "geom", name="floor", type="plane", size="1 1 .02", rgba=".18 .22 .28 1", friction="1 .01 .001")
	ET.SubElement(world, "geom", name="work_mat", type="box", pos=".235 .035 .0005", size=".155 .215 .0005", rgba=".68 .70 .67 1")
	for i, (x, y) in enumerate(FIDUCIALS):
	ET.SubElement(world, "geom", name=f"calibration_{i}", type="cylinder", pos=f"{x} {y} .0012", size=".007 .0003", rgba=".05 .95 .95 1", contype="0", conaffinity="0")
	camera_pos = np.array([.70, -.65, .68])
	target = np.array([.18, .02, .045])
	z = camera_pos - target
	z /= np.linalg.norm(z)
	x = np.cross([0., 0., 1.], z); x /= np.linalg.norm(x)
	y = np.cross(z, x)
	ET.SubElement(world, "camera", name="front", pos=" ".join(map(str, camera_pos)), xyaxes=" ".join(map(str, np.r_[x, y])), fovy="42")
	placements = [("red", .255, -.105), ("green", .285, -.025), ("blue", .19, -.075)]
	for name, x, y in placements:
	body = ET.SubElement(world, "body", name=f"block_{name}", pos=f"{x} {y} .014")
	ET.SubElement(body, "freejoint", name=f"block_{name}_free")
	ET.SubElement(body, "geom", name=f"block_{name}_solid", type="box", size=".012 .010 .012", rgba=COLORS[name], mass=".014", friction="1.5 .03 .003", condim="4", solref=".008 1")
	for sx in (-.006, .006):
	for sy in (-.005, .005):
	ET.SubElement(body, "geom", type="cylinder", pos=f"{sx} {sy} .014", size=".003 .002", rgba=COLORS[name], mass=".0001", friction="1.2 .02 .002")
	box = ET.SubElement(world, "body", name="box", pos=".20 .135 0")
	ET.SubElement(box, "geom", name="box_bottom", type="box", pos="0 0 .003", size=".062 .050 .003", rgba=".52 .36 .19 1", friction="1 .02 .002")
	for name, pos, size in [
	("left", "-.064 0 .020", ".002 .054 .020"),
	("right", ".064 0 .020", ".002 .054 .020"),
	("back", "0 .052 .020", ".062 .002 .020"),
	("front", "0 -.052 .020", ".062 .002 .020"),
	]:
	ET.SubElement(box, "geom", name=f"box_{name}", type="box", pos=pos, size=size, rgba=".66 .45 .24 1", friction="1 .02 .002")
	ET.indent(scene)
	ET.ElementTree(scene).write(OUT / "scene.xml", encoding="unicode")
	print(OUT / "scene.xml")


	if __name__ == "__main__":
	main()