Skip to content

grasp_goal

GraspGoal

Bases: SuccessCondition

GraspGoal (success condition)

Source code in OmniGibson/omnigibson/termination_conditions/grasp_goal.py
class GraspGoal(SuccessCondition):
    """
    GraspGoal (success condition)
    """

    def __init__(self, obj_name):
        self.obj_name = obj_name
        self._objs = None

        # Run super init
        super().__init__()

    def _step(self, task, env, action):
        if self._objs is None:
            self._objs = [env.scenes[i].object_registry("name", self.obj_name) for i in range(env.num_envs)]

        results = th.zeros(env.num_envs, dtype=th.bool)
        for env_idx in range(env.num_envs):
            obj = self._objs[env_idx]
            robot = env.scenes[env_idx].robots[0]
            obj_in_hand = robot._ag_obj_in_hand[robot.default_arm]
            results[env_idx] = obj_in_hand == obj
        return results