pong_catastrophe.py 文件源码

def label_and_build_mask(self, episode):
        is_catastrophe_array = np.array(
            [is_catastrophe(frame.image) for frame in episode.frames if frame.action is not None])
        # should_block_array = np.array([should_block(frame.image, frame.action) for frame in episode.frames])

        labels = np.full(len(episode.frames), fill_value=False, dtype=np.bool)
        mask = np.full(len(episode.frames), fill_value=True, dtype=np.bool)

        for i in range(len(episode.frames)):
            if i + self.block_radius + 1 >= len(episode.frames):
                mask[i] = False
                continue
            if is_catastrophe_array[i]:
                mask[i] = False
                continue
            for j in range(self.block_radius + 1):
                if is_catastrophe_array[i + j + 1]:
                    labels[i] = True
                    break
        return labels, mask