    def task_graph(self, obs):
        """
        predict the next task to transition into based on the current observation
        if the current task is not successful, stay in the current task
        else transition to the next task
        """
        success = self.compute_reward(obs)
        if success:
            return (self.task_id + 1) % 2
        return self.task_id
predict the next task to transition into based on the current observation
if the current task is not successful, stay in the current task
else transition to the next task
serl_robot_infra/franka_env/envs/wrappers.py