File size: 1,869 Bytes
393d3de |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 |
# This is a util file for various functions that retrieve object information
from xml.etree import ElementTree
from libero.envs.objects import OBJECTS_DICT, get_object_fn
EXCEPTION_DICT = {"flat_stove": "flat_stove_burner"}
def update_exception_dict(object_name, site_name):
"""Update EXCEPTION_DICT information. This is to handle some special case of affordance region naming.
Args:
object_name (str): object name
site_name (str): site name
"""
EXCEPTION_DICT[object_name] = site_name
def get_affordance_regions(objects, verbose=False):
"""_summary_
Args:
objects (MujocoObject): a dictionary of objects
verbose (bool, optional): Print additional debug information. Defaults to False.
Returns:
dict: a dictionary of object names and their affordance regions.
"""
affordances = {}
for object_name in objects.keys():
try:
obj = get_object_fn(object_name)()
# print(obj.root.findall(".//site"))
object_affordance = []
for site in obj.root.findall(".//site"):
site_name = site.get("name")
if "site" not in site_name and (
object_name not in EXCEPTION_DICT
or object_name in EXCEPTION_DICT
and site_name not in EXCEPTION_DICT[object_name]
):
# print(site_name)
# object name is already added as prefix when the object is initialized. remove them for consistency in bddl files
object_affordance.append(site_name.replace(f"{object_name}_", ""))
if len(object_affordance) > 0:
affordances[object_name] = object_affordance
except:
if verbose:
print(f"Skipping {object_name}")
return affordances
|