Files
gen_data_agent/data_gen_dependencies/action/__init__.py
2025-09-24 11:11:24 +08:00

35 lines
951 B
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 .press import PressPrismaticStage
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,
'press_prismatic': PressPrismaticStage,
"pull": PullStage,
"pull_revolute": PullRevoluteStage,
}
def build_stage(action):
if action not in ACTION_STAGE:
raise NotImplementedError
return ACTION_STAGE[action]