<commit_before>from abc import ABC, abstractmethod
import torch

from .. import utils
from babyai.agents.bot import Bot


class Agent(ABC):
    """An abstraction of the behavior of an agent. The agent is able:
    - to choose an action given an observation,
    - to analyze the feedback (i.e. reward and done state) of its action."""

    def on_reset(self):
        pass

    @abstractmethod
    def get_action(self, obs):
        pass

    @abstractmethod
    def analyze_feedback(self, reward, done):
        pass


class ModelAgent(Agent):
    """A model-based agent. This agent behaves using a model."""

    def __init__(self, model_name, obss_preprocessor, argmax):
        self.obss_preprocessor = obss_preprocessor
        self.model = utils.load_model(model_name)
        self.argmax = argmax

        if self.model.recurrent:
            self._initialize_memory()

    def _initialize_memory(self):
        self.memory = torch.zeros(1, self.model.memory_size)

    def get_action_and_outputs(self, obs):
        preprocessed_obs = self.obss_preprocessor([obs])

        with torch.no_grad():
            if self.model.recurrent:
                dist, value, self.memory = self.model(preprocessed_obs, self.memory)
            else:
                dist, value = self.model(preprocessed_obs)

        if self.argmax:
            action = dist.probs.max(1, keepdim=True)[1]
        else:
            action = dist.sample()

        return action, dist, value

    def get_action(self, obs):
        return self.get_action_and_outputs(obs)[0]

    def analyze_feedback(self, reward, done):
        if done and self.model.recurrent:
            self._initialize_memory()


class DemoAgent(Agent):
    """A demonstration-based agent. This agent behaves using demonstrations."""

    def __init__(self, demos, env_name, origin):
        self.demos_path = utils.get_demos_path(demos, env_name, origin, valid=False)
        self.demos = utils.load_demos(self.demos_path)
        self.demo_id = 0
        self.step_id = 0

    @staticmethod
    def check_obss_equality(obs1, obs2):
        if not(obs1.keys() == obs2.keys()):
            return False
        for key in obs1.keys():
            if type(obs1[key]) in (str, int):
                if not(obs1[key] == obs2[key]):
                    return False
            else:
                if not (obs1[key] == obs2[key]).all():
                    return False
        return True

    def get_action(self, obs):
        if self.demo_id >= len(self.demos):
            raise ValueError("No demonstration remaining")

        expected_obs = self.demos[self.demo_id][self.step_id][0]
        assert DemoAgent.check_obss_equality(obs, expected_obs), "The observations do not match"

        return self.demos[self.demo_id][self.step_id][1]

    def analyze_feedback(self, reward, done):
        self.step_id += 1

        if done:
            self.demo_id += 1
            self.step_id = 0


class BotAgent:
    def __init__(self, env):
        """An agent based on a GOFAI bot."""
        self.env = env
        self.on_reset()

    def on_reset(self):
        self.bot = Bot(self.env)

    def get_action(self, *args, **kwargs):
        return self.bot.step()

    def analyze_feedback(self, reward, done):
        pass
<commit_msg>make sure DemoAgent reads demos correctly (new format)<commit_after>from abc import ABC, abstractmethod
import torch

from .. import utils
from babyai.agents.bot import Bot


class Agent(ABC):
    """An abstraction of the behavior of an agent. The agent is able:
    - to choose an action given an observation,
    - to analyze the feedback (i.e. reward and done state) of its action."""

    def on_reset(self):
        pass

    @abstractmethod
    def get_action(self, obs):
        pass

    @abstractmethod
    def analyze_feedback(self, reward, done):
        pass


class ModelAgent(Agent):
    """A model-based agent. This agent behaves using a model."""

    def __init__(self, model_name, obss_preprocessor, argmax):
        self.obss_preprocessor = obss_preprocessor
        self.model = utils.load_model(model_name)
        self.argmax = argmax

        if self.model.recurrent:
            self._initialize_memory()

    def _initialize_memory(self):
        self.memory = torch.zeros(1, self.model.memory_size)

    def get_action_and_outputs(self, obs):
        preprocessed_obs = self.obss_preprocessor([obs])

        with torch.no_grad():
            if self.model.recurrent:
                dist, value, self.memory = self.model(preprocessed_obs, self.memory)
            else:
                dist, value = self.model(preprocessed_obs)

        if self.argmax:
            action = dist.probs.max(1, keepdim=True)[1]
        else:
            action = dist.sample()

        return action, dist, value

    def get_action(self, obs):
        return self.get_action_and_outputs(obs)[0]

    def analyze_feedback(self, reward, done):
        if done and self.model.recurrent:
            self._initialize_memory()


class DemoAgent(Agent):
    """A demonstration-based agent. This agent behaves using demonstrations."""

    def __init__(self, demos, env_name, origin):
        self.demos_path = utils.get_demos_path(demos, env_name, origin, valid=False)
        self.demos = utils.load_demos(self.demos_path)
        self.demo_id = 0
        self.step_id = 0

    @staticmethod
    def check_obss_equality(obs1, obs2):
        if not(obs1.keys() == obs2.keys()):
            return False
        for key in obs1.keys():
            if type(obs1[key]) in (str, int):
                if not(obs1[key] == obs2[key]):
                    return False
            else:
                if not (obs1[key] == obs2[key]).all():
                    return False
        return True

    def get_action(self, obs):
        if self.demo_id >= len(self.demos):
            raise ValueError("No demonstration remaining")
        expected_obs = self.demos[self.demo_id][0][self.step_id][0]
        assert DemoAgent.check_obss_equality(obs, expected_obs), "The observations do not match"

        return self.demos[self.demo_id][0][self.step_id][1]

    def analyze_feedback(self, reward, done):
        self.step_id += 1

        if done:
            self.demo_id += 1
            self.step_id = 0


class BotAgent:
    def __init__(self, env):
        """An agent based on a GOFAI bot."""
        self.env = env
        self.on_reset()

    def on_reset(self):
        self.bot = Bot(self.env)

    def get_action(self, *args, **kwargs):
        return self.bot.step()

    def analyze_feedback(self, reward, done):
        pass

{'fraction_non_alphanumeric': 0.10326086956521739, 'fraction_numerical': 0.006642512077294686, 'mean_word_length': 2.2911077993045206, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '770777', 'n_tokens_mistral': 2064, 'n_tokens_neox': 1967, 'n_words': 521}