<commit_before>from gym_minigrid.minigrid import *
from babyai.levels.verifier import *
from babyai.levels.verifier import (ObjDesc, pos_next_to,
                                    GoToInstr, OpenInstr, PickupInstr, PutNextInstr, BeforeInstr, AndInstr, AfterInstr)


class DisappearedBoxError(Exception):
    """
    Error that's thrown when a box is opened.
    We make the assumption that the bot cannot accomplish the mission when it happens.
    """
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


def manhattan_distance(pos, target):
    return np.abs(target[0] - pos[0]) + np.abs(target[1] - pos[1])


class Subgoal:
    """The base class for all possible Bot subgoals"""

    def __init__(self, bot=None, datum=None, reason=None):
        """
        Initializes a Subgoal object
        bot is the Bot that is trying to perform this subgoal
        datum is the primary information necessary to accomplish the subgoal
        These subgoals have extra class variables defined in theyr own __init__ method: GoNextTo
        the rest are secondary information that are used to make the bot more efficient
        """
        self.bot = bot
        self.datum = datum
        self.reason = reason

        self.update_agent_attributes()

        self.actions = self.bot.mission.actions

    def __repr__(self):
        """Mainly for debugging purposes"""
        representation = '('
        representation += type(self).__name__
        if self.datum is not None:
            representation += ': {}'.format(self.datum)
        if self.reason is not None:
            representation += ', reason: {}'.format(self.reason)
        representation += ')'
        return representation

    def update_agent_attributes(self):
        """Should be called at each step to update some elements about the agent and the environment"""
        self.pos = self.bot.mission.agent_pos
        self.dir_vec = self.bot.mission.dir_vec
        self.right_vec = self.bot.mission.right_vec
        self.fwd_pos = self.pos + self.dir_vec
        self.fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
        self.carrying = self.bot.mission.carrying

    def replan_before_action(self):
        """
        Function that updates the bot's stack given the action played.
        Should be overridden in all sub-classes.
        TODO: There are some steps that are common in both replan and get_action
        for certain subgoals, maybe the bot's speed can be improved if we do them only once
        - or at least we can factorize the code a bit

        When this function returns `True`, it means no more replanning needs to be done.
        When it returns `False`, the bot continues replanning.

        """
        raise NotImplementedError()

    def replan_after_action(self, action_taken):
        pass


    def is_exploratory(self):
        return False

    def plan_undo_action(self, action_taken):
        if action_taken == self.actions.forward:
            # check if the 'forward' action was succesful
            if not np.array_equal(self.bot.prev_agent_pos, self.pos):
                self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos))
        elif action_taken == self.actions.left:
            old_fwd_pos = self.pos + self.right_vec
            self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos + self.right_vec))
        elif action_taken == self.actions.right:
            old_fwd_pos = self.pos - self.right_vec
            self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos - self.right_vec))
        elif action_taken == self.actions.drop and self.bot.prev_carrying != self.carrying:
            # get that thing back
            assert self.fwd_cell.type in ('key', 'box', 'ball')
            self.bot.stack.append(PickupSubgoal(self.bot))
        elif action_taken == self.actions.pickup and self.bot.prev_carrying != self.carrying:
            # drop that thing where you found it
            fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
            self.bot.stack.append(DropSubgoal(self.bot))
        elif action_taken == self.actions.toggle:
            fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
            if (fwd_cell and fwd_cell.type == 'door'
                    and self.bot.fwd_door_was_open != fwd_cell.is_open):
                self.bot.stack.append(CloseSubgoal(self.bot)
                                      if fwd_cell.is_open
                                      else OpenSubgoal(self.bot))



class CloseSubgoal(Subgoal):

    def replan_after_action(self, action_taken):
        if action_taken == self.actions.toggle:
            self.bot.stack.pop()
        elif action_taken in [self.actions.forward, self.actions.left, self.actions.right]:
            self.plan_undo_action(action_taken)

    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'
        assert self.fwd_cell.is_open, 'Forward door must be open'
        return self.actions.toggle


class OpenSubgoal(Subgoal):

    def replan_after_action(self, action_taken):
        if action_taken == self.actions.toggle:
            self.bot.stack.pop()
            if self.reason == 'Unlock':
                drop_key_pos = self.bot.find_drop_pos()
                self.bot.stack.append(DropSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, drop_key_pos))
        else:
            self.plan_undo_action(action_taken)

    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:<commit_msg>a bit of commnets<commit_after>from gym_minigrid.minigrid import *
from babyai.levels.verifier import *
from babyai.levels.verifier import (ObjDesc, pos_next_to,
                                    GoToInstr, OpenInstr, PickupInstr, PutNextInstr, BeforeInstr, AndInstr, AfterInstr)


class DisappearedBoxError(Exception):
    """
    Error that's thrown when a box is opened.
    We make the assumption that the bot cannot accomplish the mission when it happens.
    """
    def __init__(self, value):
        self.value = value

    def __str__(self):
        return repr(self.value)


def manhattan_distance(pos, target):
    return np.abs(target[0] - pos[0]) + np.abs(target[1] - pos[1])


class Subgoal:
    """The base class for all possible Bot subgoals"""

    def __init__(self, bot=None, datum=None, reason=None):
        """
        Initializes a Subgoal object
        bot is the Bot that is trying to perform this subgoal
        datum is the primary information necessary to accomplish the subgoal
        These subgoals have extra class variables defined in theyr own __init__ method: GoNextTo
        the rest are secondary information that are used to make the bot more efficient
        """
        self.bot = bot
        self.datum = datum
        self.reason = reason

        self.update_agent_attributes()

        self.actions = self.bot.mission.actions

    def __repr__(self):
        """Mainly for debugging purposes"""
        representation = '('
        representation += type(self).__name__
        if self.datum is not None:
            representation += ': {}'.format(self.datum)
        if self.reason is not None:
            representation += ', reason: {}'.format(self.reason)
        representation += ')'
        return representation

    def update_agent_attributes(self):
        """Should be called at each step to update some elements about the agent and the environment"""
        self.pos = self.bot.mission.agent_pos
        self.dir_vec = self.bot.mission.dir_vec
        self.right_vec = self.bot.mission.right_vec
        self.fwd_pos = self.pos + self.dir_vec
        self.fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
        self.carrying = self.bot.mission.carrying

    def replan_before_action(self):
        """
        Function that updates the bot's stack given the action played.
        Should be overridden in all sub-classes.
        TODO: There are some steps that are common in both replan and get_action
        for certain subgoals, maybe the bot's speed can be improved if we do them only once
        - or at least we can factorize the code a bit

        When this function returns `True`, it means no more replanning needs to be done.
        When it returns `False`, the bot continues replanning.

        """
        raise NotImplementedError()

    def replan_after_action(self, action_taken):
        pass


    def is_exploratory(self):
        return False

    def plan_undo_action(self, action_taken):
        if action_taken == self.actions.forward:
            # check if the 'forward' action was succesful
            if not np.array_equal(self.bot.prev_agent_pos, self.pos):
                self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos))
        elif action_taken == self.actions.left:
            old_fwd_pos = self.pos + self.right_vec
            self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos + self.right_vec))
        elif action_taken == self.actions.right:
            old_fwd_pos = self.pos - self.right_vec
            self.bot.stack.append(GoNextToSubgoal(self.bot, self.pos - self.right_vec))
        elif action_taken == self.actions.drop and self.bot.prev_carrying != self.carrying:
            # get that thing back
            assert self.fwd_cell.type in ('key', 'box', 'ball')
            self.bot.stack.append(PickupSubgoal(self.bot))
        elif action_taken == self.actions.pickup and self.bot.prev_carrying != self.carrying:
            # drop that thing where you found it
            fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
            self.bot.stack.append(DropSubgoal(self.bot))
        elif action_taken == self.actions.toggle:
            fwd_cell = self.bot.mission.grid.get(*self.fwd_pos)
            if (fwd_cell and fwd_cell.type == 'door'
                    and self.bot.fwd_door_was_open != fwd_cell.is_open):
                self.bot.stack.append(CloseSubgoal(self.bot)
                                      if fwd_cell.is_open
                                      else OpenSubgoal(self.bot))



class CloseSubgoal(Subgoal):

    def replan_after_action(self, action_taken):
        if action_taken == self.actions.toggle:
            self.bot.stack.pop()
        elif action_taken in [self.actions.forward, self.actions.left, self.actions.right]:
            self.plan_undo_action(action_taken)

    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'
        assert self.fwd_cell.is_open, 'Forward door must be open'
        return self.actions.toggle


class OpenSubgoal(Subgoal):

    def replan_after_action(self, action_taken):
        if action_taken == self.actions.toggle:
            self.bot.stack.pop()
            if self.reason == 'Unlock':
                # The reason why this has to be planned after the action is taken
                # is because if the positions for dropping is chosen in advance,
                # then by the time the key is dropped there, it might already
                # be occupied.
                drop_key_pos = self.bot.find_drop_pos()
                self.bot.stack.append(DropSubgoal(self.bot))
                self.bot.stack.append(GoNextToSubgoal(self.bot, drop_key_pos))
        else:
            self.plan_undo_action(action_taken)

    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:
{'fraction_non_alphanumeric': 0.08724058416602613, 'fraction_numerical': 0.0006149116064565719, 'mean_word_length': 2.319979586629242, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 3, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 10, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6147612', 'n_tokens_mistral': 3847, 'n_tokens_neox': 3673, 'n_words': 1076}