|
|
import re |
|
|
|
|
|
from libero.envs.objects import OBJECTS_DICT |
|
|
from libero.utils.bddl_generation_utils import ( |
|
|
get_affordance_region_kwargs_list_from_fixture_info, |
|
|
get_object_dict, |
|
|
) |
|
|
from libero.utils.object_utils import get_affordance_regions |
|
|
|
|
|
MU_DICT = {} |
|
|
|
|
|
SCENE_DICT = {} |
|
|
|
|
|
|
|
|
def register_mu(scene_type="general"): |
|
|
def _func(target_class): |
|
|
"""For reusing initial conditions easily, we register each pre-defined initial conditions in a dictionary.""" |
|
|
key = "_".join( |
|
|
re.sub(r"([A-Z])", r" \1", target_class.__name__).split() |
|
|
).lower() |
|
|
MU_DICT[key.lower()] = target_class |
|
|
if scene_type not in SCENE_DICT: |
|
|
SCENE_DICT[scene_type.lower()] = [] |
|
|
SCENE_DICT[scene_type.lower()].append(target_class) |
|
|
return target_class |
|
|
|
|
|
return _func |
|
|
|
|
|
|
|
|
def get_scene_dict(scene_type=None): |
|
|
if scene_type is None: |
|
|
return SCENE_DICT |
|
|
else: |
|
|
return SCENE_DICT[scene_type.lower()] |
|
|
|
|
|
|
|
|
def get_scene_class(scene_name): |
|
|
return MU_DICT[scene_name.lower()] |
|
|
|
|
|
|
|
|
class InitialSceneTemplates: |
|
|
def __init__( |
|
|
self, workspace_name="main_table", fixture_num_info={}, object_num_info={} |
|
|
): |
|
|
|
|
|
self.workspace_name = workspace_name |
|
|
|
|
|
|
|
|
self.fixture_object_dict = get_object_dict(fixture_num_info) |
|
|
self.movable_object_dict = get_object_dict(object_num_info) |
|
|
|
|
|
affordances = get_affordance_regions(OBJECTS_DICT) |
|
|
affordance_fixture_info_dict = {} |
|
|
for fixture_category_name in self.fixture_object_dict.keys(): |
|
|
if ( |
|
|
fixture_category_name != self.workspace_name |
|
|
and fixture_category_name != "table" |
|
|
and fixture_category_name != "living_room_table" |
|
|
and fixture_category_name != "study_table" |
|
|
and fixture_category_name != "kitchen_table" |
|
|
): |
|
|
for fixture_name in self.fixture_object_dict[fixture_category_name]: |
|
|
affordance_fixture_info_dict[fixture_name] = affordances[ |
|
|
fixture_category_name |
|
|
] |
|
|
for category_name in self.movable_object_dict.keys(): |
|
|
if category_name in affordances: |
|
|
for object_name in self.movable_object_dict[category_name]: |
|
|
affordance_fixture_info_dict[object_name] = affordances[ |
|
|
category_name |
|
|
] |
|
|
|
|
|
self.affordance_region_kwargs_list = ( |
|
|
get_affordance_region_kwargs_list_from_fixture_info( |
|
|
affordance_fixture_info_dict |
|
|
) |
|
|
) |
|
|
|
|
|
self.regions = {} |
|
|
self.define_regions() |
|
|
|
|
|
@property |
|
|
def possible_objects_of_interest(self): |
|
|
|
|
|
|
|
|
objects_of_interest = [] |
|
|
for category_name in self.fixture_object_dict.keys(): |
|
|
objects_of_interest += self.fixture_object_dict[category_name] |
|
|
for category_name in self.movable_object_dict.keys(): |
|
|
objects_of_interest += self.movable_object_dict[category_name] |
|
|
return objects_of_interest |
|
|
|
|
|
@property |
|
|
def movable_objects(self): |
|
|
return list(self.movable_object_dict) |
|
|
|
|
|
def define_regions(self): |
|
|
"""Override this method to define the layout of a scene.""" |
|
|
raise NotImplementedError |
|
|
|
|
|
def get_region_dict( |
|
|
self, |
|
|
region_centroid_xy, |
|
|
region_name, |
|
|
target_name=None, |
|
|
region_half_len=0.02, |
|
|
yaw_rotation=(0.0, 0.0), |
|
|
): |
|
|
"""This is a function that creates a default region with rectangular shape.""" |
|
|
if target_name is None: |
|
|
target_name = self.workspace_name |
|
|
region_key_value = { |
|
|
region_name: { |
|
|
"target": target_name, |
|
|
"ranges": [ |
|
|
( |
|
|
region_centroid_xy[0] - region_half_len, |
|
|
region_centroid_xy[1] - region_half_len, |
|
|
region_centroid_xy[0] + region_half_len, |
|
|
region_centroid_xy[1] + region_half_len, |
|
|
) |
|
|
], |
|
|
"yaw_rotation": [yaw_rotation], |
|
|
} |
|
|
} |
|
|
print |
|
|
return region_key_value |
|
|
|
|
|
@property |
|
|
def init_states(self): |
|
|
raise NotImplementedError |
|
|
|
