def _get_shaped_reward()

in mae_envs/envs/box_locking.py [0:0]

• 18 lines of code
• 12 McCabe index (conditional complexity)

    def _get_shaped_reward(self, new_next_obj, new_next_obj_dist, new_spawn_pos_dist):
        '''
            Calculates the shaped reward based on the change in distance from the target
        '''
        rew = 0
        if (self.next_obj is not None) and (new_next_obj == self.next_obj):
            rew += (self.next_obj_dist - new_next_obj_dist) * self.shaped_reward_scale
        elif ((self.next_obj is not None) and (new_next_obj != self.next_obj)):
            if self.objs_locked[self.next_obj] == 1:
                # previous target object locked
                rew += self.next_obj_dist * self.shaped_reward_scale
            else:
                # previously locked object unlocked
                rew -= new_next_obj_dist * self.shaped_reward_scale
        elif (self.next_obj is None) and (new_next_obj is not None):
            # previously locked object unlocked
            rew -= new_next_obj_dist * self.shaped_reward_scale
        elif (self.next_obj is None) and (new_next_obj is None):
            if self.need_return:
                # all objects locked; agent is rewarded for returning to its spawning point
                rew += (self.spawn_pos_dist - new_spawn_pos_dist) * self.shaped_reward_scale

        return rew