def step()

in bisk/tasks/goalwall.py [0:0]

• 38 lines of code
• 11 McCabe index (conditional complexity)

    def step(self, action):
        self.ball_yz = None
        btbefore = self.ball_touched
        obs, reward, done, info = super().step(action)

        goal_hit = None
        goal_dists = []
        goal_sizes = []
        if self.ball_yz is not None:
            if self.is_2d:
                goals = ('goal', )
            else:
                goals = ('goal1', 'goal2')
            for g in goals:
                d = np.linalg.norm(self.ball_yz -
                                   self.p.named.data.site_xpos[g][1:3])
                goal_dists.append(d)
                goal_sizes.append(self.p.named.model.site_size[g][2])
                if d <= self.p.named.model.site_size[g][2]:
                    goal_hit = g
                    break

        score = 0
        if goal_hit == 'goal' or goal_hit == 'goal1':
            score = 1
        elif goal_hit == 'goal2':
            score = 2
        info['score'] = score
        reward = score
        if self.touch_ball_reward > 0 and self.ball_touched != btbefore:
            reward += self.touch_ball_reward

        # Zero reward if we're beyond the line
        lpos = self.p.named.data.site_xpos['line', 'x']
        if self.robot_pos[0] > lpos:
            reward = 0

        # Once we've hit the wall we're done
        if self.ball_yz is not None:
            done = True

        if info.get('fell_over', False):
            reward = -1
            done = True
        return obs, reward, done, info