def replan_before_action()

in data/envs/babyai/bot_agent.py [0:0]

• 33 lines of code
• 15 McCabe index (conditional complexity)

    def replan_before_action(self):
        assert self.fwd_cell is not None, "Forward cell is empty"
        assert self.fwd_cell.type == "door", "Forward cell has to be a door"

        # If the door is locked, go find the key and then return
        # TODO: do we really need to be in front of the locked door
        # to realize that we need the key for it ?
        got_the_key = self.carrying and self.carrying.type == "key" and self.carrying.color == self.fwd_cell.color
        if self.fwd_cell.is_locked and not got_the_key:
            # Find the key
            key_desc = ObjDesc("key", self.fwd_cell.color)
            key_desc.find_matching_objs(self.bot.mission)

            # If we're already carrying something
            if self.carrying:
                self.bot.stack.pop()

                # Find a location to drop what we're already carrying
                drop_pos_cur = self.bot._find_drop_pos()

                # Take back the object being carried
                self.bot.stack.append(PickupSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, drop_pos_cur))

                # Go back to the door and open it
                self.bot.stack.append(OpenSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, tuple(self.fwd_pos)))

                # Go to the key and pick it up
                self.bot.stack.append(PickupSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, key_desc))

                # Drop the object being carried
                self.bot.stack.append(DropSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, drop_pos_cur))
            else:
                # This branch is will be used very rarely, given that
                # GoNextToSubGoal(..., reason='Open') should plan
                # going to the key before we get to stand right in front of a door.
                # But the agent can be spawned right in front of a open door,
                # for which we case we do need this code.

                self.bot.stack.pop()

                # Go back to the door and open it
                self.bot.stack.append(OpenSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, tuple(self.fwd_pos)))

                # Go to the key and pick it up
                self.bot.stack.append(PickupSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, key_desc))
            return

        if self.fwd_cell.is_open:
            self.bot.stack.append(CloseSubgoal(self.bot))
            return

        if self.fwd_cell.is_locked and self.reason is None:
            self.bot.stack.pop()
            self.bot.stack.append(OpenSubgoal(self.bot, reason="Unlock"))
            return

        return self.actions.toggle