File size: 3,997 Bytes
e479c46 | 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 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 86 87 | import gymnasium as gym
import mani_skill2_real2sim.envs
ENVIRONMENTS = [
"google_robot_pick_coke_can",
"google_robot_pick_horizontal_coke_can",
"google_robot_pick_vertical_coke_can",
"google_robot_pick_standing_coke_can",
"google_robot_pick_object",
"google_robot_move_near_v0",
"google_robot_move_near_v1",
"google_robot_move_near",
"google_robot_open_drawer",
"google_robot_open_top_drawer",
"google_robot_open_middle_drawer",
"google_robot_open_bottom_drawer",
"google_robot_close_drawer",
"google_robot_close_top_drawer",
"google_robot_close_middle_drawer",
"google_robot_close_bottom_drawer",
"google_robot_place_in_closed_drawer",
"google_robot_place_in_closed_top_drawer",
"google_robot_place_in_closed_middle_drawer",
"google_robot_place_in_closed_bottom_drawer",
"google_robot_place_apple_in_closed_top_drawer",
"widowx_spoon_on_towel",
"widowx_carrot_on_plate",
"widowx_stack_cube",
"widowx_put_eggplant_in_basket",
"widowx_put_eggplant_in_sink", # TODO(YL): Add this task
"widowx_open_drawer",
"widowx_close_drawer",
]
ENVIRONMENT_MAP = {
"google_robot_pick_coke_can": ("GraspSingleOpenedCokeCanInScene-v0", {}),
"google_robot_pick_horizontal_coke_can": (
"GraspSingleOpenedCokeCanInScene-v0",
{"lr_switch": True},
),
"google_robot_pick_vertical_coke_can": (
"GraspSingleOpenedCokeCanInScene-v0",
{"laid_vertically": True},
),
"google_robot_pick_standing_coke_can": (
"GraspSingleOpenedCokeCanInScene-v0",
{"upright": True},
),
"google_robot_pick_object": ("GraspSingleRandomObjectInScene-v0", {}),
"google_robot_move_near": ("MoveNearGoogleBakedTexInScene-v1", {}),
"google_robot_move_near_v0": ("MoveNearGoogleBakedTexInScene-v0", {}),
"google_robot_move_near_v1": ("MoveNearGoogleBakedTexInScene-v1", {}),
"google_robot_open_drawer": ("OpenDrawerCustomInScene-v0", {}),
"google_robot_open_top_drawer": ("OpenTopDrawerCustomInScene-v0", {}),
"google_robot_open_middle_drawer": ("OpenMiddleDrawerCustomInScene-v0", {}),
"google_robot_open_bottom_drawer": ("OpenBottomDrawerCustomInScene-v0", {}),
"google_robot_close_drawer": ("CloseDrawerCustomInScene-v0", {}),
"google_robot_close_top_drawer": ("CloseTopDrawerCustomInScene-v0", {}),
"google_robot_close_middle_drawer": ("CloseMiddleDrawerCustomInScene-v0", {}),
"google_robot_close_bottom_drawer": ("CloseBottomDrawerCustomInScene-v0", {}),
"google_robot_place_in_closed_drawer": ("PlaceIntoClosedDrawerCustomInScene-v0", {}),
"google_robot_place_in_closed_top_drawer": ("PlaceIntoClosedTopDrawerCustomInScene-v0", {}),
"google_robot_place_in_closed_middle_drawer": ("PlaceIntoClosedMiddleDrawerCustomInScene-v0", {}),
"google_robot_place_in_closed_bottom_drawer": ("PlaceIntoClosedBottomDrawerCustomInScene-v0", {}),
"google_robot_place_apple_in_closed_top_drawer": (
"PlaceIntoClosedTopDrawerCustomInScene-v0",
{"model_ids": "baked_apple_v2"}
),
"widowx_spoon_on_towel": ("PutSpoonOnTableClothInScene-v0", {}),
"widowx_carrot_on_plate": ("PutCarrotOnPlateInScene-v0", {}),
"widowx_stack_cube": ("StackGreenCubeOnYellowCubeBakedTexInScene-v0", {}),
"widowx_put_eggplant_in_basket": ("PutEggplantInBasketScene-v0", {}),
"widowx_put_eggplant_in_sink": ("PutEggplantInSinkScene-v0", {}),
"widowx_open_drawer": ("OpenSmallDrawerCustomInScene-v0", {}),
"widowx_close_drawer": ("CloseSmallDrawerCustomInScene-v0", {}),
}
def make(task_name):
"""Creates simulated eval environment from task name."""
assert task_name in ENVIRONMENTS, f"Task {task_name} is not supported. Environments: \n {ENVIRONMENTS}"
env_name, kwargs = ENVIRONMENT_MAP[task_name]
kwargs["prepackaged_config"] = True
env = gym.make(env_name, obs_mode="rgbd", **kwargs)
print(f" \n\n - Created environment {env_name}")
return env
|