Skip to content

grasp_reward

GraspReward

Bases: BaseRewardFunction

A composite reward function for grasping tasks. This reward function not only evaluates the success of object grasping but also considers various penalties and efficiencies.

The reward is calculated based on several factors: - Grasping reward: A positive reward is given if the robot is currently grasping the specified object. - Distance reward: A reward based on the inverse exponential distance between the end-effector and the object. - Regularization penalty: Penalizes large magnitude actions to encourage smoother and more energy-efficient movements. - Position and orientation penalties: Discourages excessive movement of the end-effector. - Collision penalty: Penalizes collisions with the environment or other objects.

Attributes:

Name Type Description
obj_name str

Name of the object to grasp.
dist_coeff float

Coefficient for the distance reward calculation.
grasp_reward float

Reward given for successfully grasping the object.
collision_penalty float

Penalty incurred for any collision.
eef_position_penalty_coef float

Coefficient for the penalty based on end-effector's position change.
eef_orientation_penalty_coef float

Coefficient for the penalty based on end-effector's orientation change.
regularization_coef float

Coefficient for penalizing large actions.
Source code in OmniGibson/omnigibson/reward_functions/grasp_reward.py 
 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
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
class GraspReward(BaseRewardFunction):
    """
    A composite reward function for grasping tasks. This reward function not only evaluates the success of object grasping
    but also considers various penalties and efficiencies.

    The reward is calculated based on several factors:
    - Grasping reward: A positive reward is given if the robot is currently grasping the specified object.
    - Distance reward: A reward based on the inverse exponential distance between the end-effector and the object.
    - Regularization penalty: Penalizes large magnitude actions to encourage smoother and more energy-efficient movements.
    - Position and orientation penalties: Discourages excessive movement of the end-effector.
    - Collision penalty: Penalizes collisions with the environment or other objects.

    Attributes:
        obj_name (str): Name of the object to grasp.
        dist_coeff (float): Coefficient for the distance reward calculation.
        grasp_reward (float): Reward given for successfully grasping the object.
        collision_penalty (float): Penalty incurred for any collision.
        eef_position_penalty_coef (float): Coefficient for the penalty based on end-effector's position change.
        eef_orientation_penalty_coef (float): Coefficient for the penalty based on end-effector's orientation change.
        regularization_coef (float): Coefficient for penalizing large actions.
    """

    def __init__(
        self,
        obj_name,
        dist_coeff,
        grasp_reward,
        collision_penalty,
        eef_position_penalty_coef,
        eef_orientation_penalty_coef,
        regularization_coef,
    ):
        # Store internal vars
        self.prev_grasping = None
        self.prev_eef_pos = None
        self.prev_eef_quat = None
        self.obj_name = obj_name
        self._objs = None  # list of objects per scene, populated lazily
        self.dist_coeff = dist_coeff
        self.grasp_reward = grasp_reward
        self.collision_penalty = collision_penalty
        self.eef_position_penalty_coef = eef_position_penalty_coef
        self.eef_orientation_penalty_coef = eef_orientation_penalty_coef
        self.regularization_coef = regularization_coef

        # Run super
        super().__init__()

    def _step(self, task, env, action):
        # Lazily populate per-scene object references
        if self._objs is None:
            self._objs = [env.scenes[i].object_registry("name", self.obj_name) for i in range(env.num_envs)]

        rewards = th.zeros(env.num_envs, dtype=th.float32)
        infos = [dict() for _ in range(env.num_envs)]

        for env_idx in range(env.num_envs):
            obj = self._objs[env_idx]
            scene = env.scenes[env_idx]
            robot_idx = 0
            robot = scene.robots[robot_idx]
            obj_in_hand = robot._ag_obj_in_hand[robot.default_arm]
            current_grasping = obj_in_hand == obj

            info = {"grasp_success": current_grasping}

            # Reward varying based on combination of whether the robot was previously grasping the desired object
            # and is currently grasping the desired object
            reward = 0.0

            # Penalize large actions. action is shaped (num_envs, action_dim), so we must select
            # this env's row and, within it, only the slice belonging to this robot (keyed by
            # scene + robot, matching the action layout in Environment.step).
            robot_action_start = sum(r.action_dim for r in scene.robots[:robot_idx])
            robot_action = action[env_idx][robot_action_start : robot_action_start + robot.action_dim]
            action_mag = th.sum(th.abs(robot_action))
            regularization_penalty = -(action_mag * self.regularization_coef)
            reward += regularization_penalty
            info["regularization_penalty_factor"] = action_mag
            info["regularization_penalty"] = regularization_penalty

            # Penalize based on the magnitude of the action
            eef_pos = robot.get_eef_position(robot.default_arm)
            info["position_penalty_factor"] = 0.0
            info["position_penalty"] = 0.0
            if self.prev_eef_pos is not None and self.prev_eef_pos[env_idx] is not None:
                eef_pos_delta = T.l2_distance(self.prev_eef_pos[env_idx], eef_pos)
                position_penalty = -eef_pos_delta * self.eef_position_penalty_coef
                reward += position_penalty
                info["position_penalty_factor"] = eef_pos_delta
                info["position_penalty"] = position_penalty

            eef_quat = robot.get_eef_orientation(robot.default_arm)
            info["rotation_penalty_factor"] = 0.0
            info["rotation_penalty"] = 0.0
            if self.prev_eef_quat is not None and self.prev_eef_quat[env_idx] is not None:
                delta_rot = T.get_orientation_diff_in_radian(self.prev_eef_quat[env_idx], eef_quat)
                rotation_penalty = -delta_rot * self.eef_orientation_penalty_coef
                reward += rotation_penalty
                info["rotation_penalty_factor"] = delta_rot.item()
                info["rotation_penalty"] = rotation_penalty.item()

            # Penalize robot for colliding with an object
            info["collision_penalty_factor"] = 0.0
            info["collision_penalty"] = 0.0
            if detect_robot_collision_in_sim(robot, filter_objs=[obj]):
                reward += -self.collision_penalty
                info["collision_penalty_factor"] = 1.0
                info["collision_penalty"] = -self.collision_penalty

            # If we're not currently grasping
            info["grasp_reward_factor"] = 0.0
            info["grasp_reward"] = 0.0
            info["pregrasp_dist"] = 0.0
            info["pregrasp_dist_reward_factor"] = 0.0
            info["pregrasp_dist_reward"] = 0.0
            info["postgrasp_dist"] = 0.0
            info["postgrasp_dist_reward_factor"] = 0.0
            info["postgrasp_dist_reward"] = 0.0
            if not current_grasping:
                # TODO: If we dropped the object recently, penalize for that
                obj_center = obj.get_position_orientation()[0]
                dist = T.l2_distance(eef_pos, obj_center)
                dist_reward = math.exp(-dist) * self.dist_coeff
                reward += dist_reward
                info["pregrasp_dist"] = dist
                info["pregrasp_dist_reward_factor"] = math.exp(-dist)
                info["pregrasp_dist_reward"] = dist_reward
            else:
                # We are currently grasping - first apply a grasp reward
                reward += self.grasp_reward
                info["grasp_reward_factor"] = 1.0
                info["grasp_reward"] = self.grasp_reward

                # Then apply a distance reward to take us to a tucked position
                robot_center = robot.links["torso_lift_link"].get_position_orientation()[0]
                obj_center = obj.get_position_orientation()[0]
                dist = T.l2_distance(robot_center, obj_center)
                dist_reward = math.exp(-dist) * self.dist_coeff
                reward += dist_reward
                info["postgrasp_dist"] = dist
                info["postgrasp_dist_reward_factor"] = math.exp(-dist)
                info["postgrasp_dist_reward"] = dist_reward

            rewards[env_idx] = reward
            infos[env_idx] = info

            # Update per-env tracking
            if self.prev_eef_pos is not None:
                self.prev_eef_pos[env_idx] = eef_pos
            if self.prev_eef_quat is not None:
                self.prev_eef_quat[env_idx] = eef_quat
            if self.prev_grasping is not None:
                self.prev_grasping[env_idx] = current_grasping

        return rewards, infos

    def reset(self, task, env, env_indices):
        """
        Reward function-specific reset

        Args:
            task (BaseTask): Task instance
            env (Environment): Environment instance
            env_indices (list): List of environment indices to reset
        """
        super().reset(task, env, env_indices)
        if self.prev_grasping is None:
            self.prev_grasping = [False] * env.num_envs
            self.prev_eef_pos = [None] * env.num_envs
            self.prev_eef_quat = [None] * env.num_envs
        for idx in env_indices:
            self.prev_grasping[idx] = False
            self.prev_eef_pos[idx] = None
            self.prev_eef_quat[idx] = None

reset(task, env, env_indices)

Reward function-specific reset

Parameters:

Name Type Description Default
task BaseTask

Task instance

 required
env Environment

Environment instance

 required
env_indices list

List of environment indices to reset

 required
Source code in OmniGibson/omnigibson/reward_functions/grasp_reward.py 
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
def reset(self, task, env, env_indices):
    """
    Reward function-specific reset

    Args:
        task (BaseTask): Task instance
        env (Environment): Environment instance
        env_indices (list): List of environment indices to reset
    """
    super().reset(task, env, env_indices)
    if self.prev_grasping is None:
        self.prev_grasping = [False] * env.num_envs
        self.prev_eef_pos = [None] * env.num_envs
        self.prev_eef_quat = [None] * env.num_envs
    for idx in env_indices:
        self.prev_grasping[idx] = False
        self.prev_eef_pos[idx] = None
        self.prev_eef_quat[idx] = None