You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
78 lines
2.3 KiB
78 lines
2.3 KiB
import h5py
|
|
import json
|
|
import argparse
|
|
import os
|
|
from shutil import copyfile
|
|
import robosuite
|
|
import xml.etree.ElementTree as ET
|
|
from tqdm import tqdm
|
|
import numpy as np
|
|
|
|
|
|
def get_eef_panda_site(prefix):
|
|
"""
|
|
Helper function for new element added in robosuite v1.5
|
|
"""
|
|
|
|
eef_panda_site = (
|
|
"""<site name="right_center" pos="0 0 0" size="0.01" rgba="1 0.3 0.3 1" group="2"/> """
|
|
)
|
|
|
|
tree = ET.ElementTree(ET.fromstring(eef_panda_site))
|
|
root = tree.getroot()
|
|
for elem in root.iter():
|
|
elem.attrib["name"] = "{}_{}".format(prefix, elem.get("name"))
|
|
|
|
return tree
|
|
|
|
|
|
def update_gripper_name(gripper_name: str):
|
|
tokens = gripper_name.split("_")
|
|
tokens.insert(1, "right")
|
|
return "_".join(tokens)
|
|
|
|
|
|
def convert_xml(xml_str):
|
|
"""
|
|
Postprocess xml string generated by robosuite to be compatible with robosuite v1.5
|
|
This script should not the xml string if it was already generated using robosuite v1.5
|
|
Args:
|
|
xml_str (str): xml string to process (from robosuite v1.4.1)
|
|
"""
|
|
|
|
root = ET.fromstring(xml_str)
|
|
worldbody = root.find("worldbody")
|
|
bodies_to_process = []
|
|
|
|
for body in worldbody.iter("body"):
|
|
body_name = body.get("name")
|
|
if "link0" in body_name:
|
|
if body not in bodies_to_process:
|
|
bodies_to_process.append(body)
|
|
|
|
# add elements introduced in v1.5
|
|
for body in bodies_to_process:
|
|
prefix = body.get("name").split("_")[0]
|
|
eef_tree = get_eef_panda_site(prefix)
|
|
body.append(eef_tree.getroot())
|
|
|
|
assets = root.find("asset")
|
|
|
|
# change path from mounts to bases
|
|
for asset in assets.iter("mesh"):
|
|
path = asset.get("file")
|
|
assert "robosuite/models/assets" in path or "robocasa/models/assets" in path, print(path)
|
|
if "mounts" in path:
|
|
assert not "robot" in path
|
|
new_path = path.replace(
|
|
"robosuite/models/assets/mounts", "robosuite/models/assets/bases"
|
|
)
|
|
asset.set("file", new_path)
|
|
|
|
# update naming prefixes
|
|
for elem in root.iter():
|
|
for attr, value in elem.attrib.items():
|
|
if value.startswith("mount"):
|
|
elem.set(attr, value.replace("mount", "fixed_base", 1))
|
|
elif value.startswith("gripper"):
|
|
elem.set(attr, update_gripper_name(value))
|