# Path: babyai/rl/utils/supervised_losses.py
# class ExtraInfoCollector:
# class SupervisedLossUpdater:
#     def __init__(self, aux_info, shape, device):
#     def process(self, env_info):
#     def fill_dictionaries(self, index, env_info, extra_predictions):
#     def end_collection(self, exps):
#     def __init__(self, aux_info, supervised_loss_coef, recurrence, device):
#     def init_epoch(self):
#     def init_batch(self):
#     def eval_subbatch(self, extra_predictions, sb):
#     def update_batch_values(self):
#     def update_epoch_logs(self):
#     def end_training(self, logs):


import torch
import torch.nn as nn
import torch.nn.functional as F
import babyai.rl
from torch.autograd import Variable
from torch.distributions.categorical import Categorical
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from babyai.rl.utils.supervised_losses import required_heads
            nn.Tanh(),
            nn.Linear(64, action_space.n)
        )

        # Define critic's model
        self.critic = nn.Sequential(
            nn.Linear(self.embedding_size, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )

        # Initialize parameters correctly
        self.apply(initialize_parameters)

        # Define head for extra info
        if self.aux_info:
            self.extra_heads = None
            self.add_heads()

    def add_heads(self):
        '''
        When using auxiliary tasks, the environment yields at each step some binary, continous, or multiclass
        information. The agent needs to predict those information. This function add extra heads to the model
        that output the predictions. There is a head per extra information (the head type depends on the extra
        information type).
        '''
        self.extra_heads = nn.ModuleDict()
        for info in self.aux_info:
            if required_heads[info] == 'binary':
                self.extra_heads[info] = nn.Linear(self.embedding_size, 1)

# Path: babyai/rl/utils/supervised_losses.py
# class ExtraInfoCollector:
# class SupervisedLossUpdater:
#     def __init__(self, aux_info, shape, device):
#     def process(self, env_info):
#     def fill_dictionaries(self, index, env_info, extra_predictions):
#     def end_collection(self, exps):
#     def __init__(self, aux_info, supervised_loss_coef, recurrence, device):
#     def init_epoch(self):
#     def init_batch(self):
#     def eval_subbatch(self, extra_predictions, sb):
#     def update_batch_values(self):
#     def update_epoch_logs(self):
#     def end_training(self, logs):


# Path: babyai/model.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import babyai.rl
from torch.autograd import Variable
from torch.distributions.categorical import Categorical
from torch.nn.utils.rnn import pack_padded_sequence, pad_packed_sequence
from babyai.rl.utils.supervised_losses import required_heads

            nn.Tanh(),
            nn.Linear(64, action_space.n)
        )

        # Define critic's model
        self.critic = nn.Sequential(
            nn.Linear(self.embedding_size, 64),
            nn.Tanh(),
            nn.Linear(64, 1)
        )

        # Initialize parameters correctly
        self.apply(initialize_parameters)

        # Define head for extra info
        if self.aux_info:
            self.extra_heads = None
            self.add_heads()

    def add_heads(self):
        '''
        When using auxiliary tasks, the environment yields at each step some binary, continous, or multiclass
        information. The agent needs to predict those information. This function add extra heads to the model
        that output the predictions. There is a head per extra information (the head type depends on the extra
        information type).
        '''
        self.extra_heads = nn.ModuleDict()
        for info in self.aux_info:
            if required_heads[info] == 'binary':
                self.extra_heads[info] = nn.Linear(self.embedding_size, 1)

            elif required_heads[info].startswith('multiclass'):