Files
gen_data_agent/data_gen_dependencies/action/__init__.py

42 lines
1.2 KiB
Python

from .grasp import GraspStage, PickStage, HookStage
from .place import PlaceStage
from .insert import InsertStage, HitStage
from .slide import SlideStage
from .pour import PourStage
from .pull import PullStage
from .push import PushStage
from .pull_revolute import PullRevoluteStage
from .push_revolute import PushRevoluteStage
from .press import PressPrismaticStage
from .twist import TwistStage
ACTION_STAGE = {
"grasp": GraspStage,
"pick": PickStage,
"hook": HookStage,
"place": PlaceStage,
"insert": InsertStage,
"slide": SlideStage,
"shave": NotImplemented,
"brush": NotImplemented,
"wipe": NotImplemented,
"hit": NotImplemented,
"pour": PourStage,
"push": PushStage,
'click': PressPrismaticStage,
'touch': PressPrismaticStage,
"pull": PullStage,
# ---- Articulate ----
"pull_revolute": PullRevoluteStage,
"push_revolute": PushRevoluteStage,
"pull_prismatic":PullStage,
"push_prismatic":PushStage,
"twist": TwistStage,
'press_prismatic': PressPrismaticStage,
}
def build_stage(action):
if action not in ACTION_STAGE:
raise NotImplementedError
return ACTION_STAGE[action]