# 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
        self.bn2 = nn.BatchNorm2d(out_features)

        self.weight = nn.Linear(in_features, out_features)
        self.bias = nn.Linear(in_features, out_features)

        self.apply(initialize_parameters)

    def forward(self, x, y):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)
        weight = self.weight(y).unsqueeze(2).unsqueeze(3)
        bias = self.bias(y).unsqueeze(2).unsqueeze(3)
        out = x * weight + bias
        return F.relu(self.bn2(out))


class ImageBOWEmbedding(nn.Module):
   def __init__(self, max_value, embedding_dim):
       super().__init__()
       self.max_value = max_value
       self.embedding_dim = embedding_dim
       self.embedding = nn.Embedding(3 * max_value, embedding_dim)
       self.apply(initialize_parameters)

   def forward(self, inputs):
       offsets = torch.Tensor([0, self.max_value, 2 * self.max_value]).to(inputs.device)
       inputs = (inputs + offsets[None, :, None, None]).long()
       return self.embedding(inputs).sum(1).permute(0, 3, 1, 2)



# 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

        self.bn2 = nn.BatchNorm2d(out_features)

        self.weight = nn.Linear(in_features, out_features)
        self.bias = nn.Linear(in_features, out_features)

        self.apply(initialize_parameters)

    def forward(self, x, y):
        x = F.relu(self.bn1(self.conv1(x)))
        x = self.conv2(x)
        weight = self.weight(y).unsqueeze(2).unsqueeze(3)
        bias = self.bias(y).unsqueeze(2).unsqueeze(3)
        out = x * weight + bias
        return F.relu(self.bn2(out))


class ImageBOWEmbedding(nn.Module):
   def __init__(self, max_value, embedding_dim):
       super().__init__()
       self.max_value = max_value
       self.embedding_dim = embedding_dim
       self.embedding = nn.Embedding(3 * max_value, embedding_dim)
       self.apply(initialize_parameters)

   def forward(self, inputs):
       offsets = torch.Tensor([0, self.max_value, 2 * self.max_value]).to(inputs.device)
       inputs = (inputs + offsets[None, :, None, None]).long()
       return self.embedding(inputs).sum(1).permute(0, 3, 1, 2)



class ACModel(nn.Module, babyai.rl.RecurrentACModel):