# Path: babyai/rl/utils/dictlist.py
# class DictList(dict):
#     """A dictionnary of lists of same size. Dictionnary items can be
#     accessed using `.` notation and list items using `[]` notation.
# 
#     Example:
#         >>> d = DictList({"a": [[1, 2], [3, 4]], "b": [[5], [6]]})
#         >>> d.a
#         [[1, 2], [3, 4]]
#         >>> d[0]
#         DictList({"a": [1, 2], "b": [5]})
#     """
# 
#     __getattr__ = dict.__getitem__
#     __setattr__ = dict.__setitem__
# 
#     def __len__(self):
#         return len(next(iter(dict.values(self))))
# 
#     def __getitem__(self, index):
#         return DictList({key: value[index] for key, value in dict.items(self)})
# 
#     def __setitem__(self, index, d):
#         for key, value in d.items():
#             dict.__getitem__(self, key)[index] = value


import torch
import torch.nn.functional as F
import numpy
from babyai.rl.utils import DictList
                raise ValueError("{} not supported".format(required_heads[info]))

    def process(self, env_info):
        # env_info is now a tuple of dicts
        env_info = [{k: v for k, v in dic.items() if k in self.aux_info} for dic in env_info]
        env_info = {k: [env_info[_][k] for _ in range(len(env_info))] for k in env_info[0].keys()}
        # env_info is now a dict of lists
        return env_info

    def fill_dictionaries(self, index, env_info, extra_predictions):
        for info in self.aux_info:
            dtype = torch.long if required_heads[info].startswith('multiclass') else torch.float
            self.collected_info[info][index] = torch.tensor(env_info[info], dtype=dtype, device=self.device)
            self.extra_predictions[info][index] = extra_predictions[info]

    def end_collection(self, exps):
        collected_info = dict()
        extra_predictions = dict()
        for info in self.aux_info:
            # T x P -> P x T -> P * T
            collected_info[info] = self.collected_info[info].transpose(0, 1).reshape(-1)
            if required_heads[info] == 'binary' or required_heads[info].startswith('continuous'):
                # T x P x 1 -> P x T x 1 -> P * T
                extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1)
            elif type(required_heads[info]) == int:
                # T x P x k -> P x T x k -> (P * T) x k
                k = required_heads[info]  # number of classes
                extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1, k)
        # convert the dicts to DictLists, and add them to the exps DictList.
        exps.collected_info = DictList(collected_info)

# Path: babyai/rl/utils/dictlist.py
# class DictList(dict):
#     """A dictionnary of lists of same size. Dictionnary items can be
#     accessed using `.` notation and list items using `[]` notation.
# 
#     Example:
#         >>> d = DictList({"a": [[1, 2], [3, 4]], "b": [[5], [6]]})
#         >>> d.a
#         [[1, 2], [3, 4]]
#         >>> d[0]
#         DictList({"a": [1, 2], "b": [5]})
#     """
# 
#     __getattr__ = dict.__getitem__
#     __setattr__ = dict.__setitem__
# 
#     def __len__(self):
#         return len(next(iter(dict.values(self))))
# 
#     def __getitem__(self, index):
#         return DictList({key: value[index] for key, value in dict.items(self)})
# 
#     def __setitem__(self, index, d):
#         for key, value in d.items():
#             dict.__getitem__(self, key)[index] = value


# Path: babyai/rl/utils/supervised_losses.py
import torch
import torch.nn.functional as F
import numpy
from babyai.rl.utils import DictList

                raise ValueError("{} not supported".format(required_heads[info]))

    def process(self, env_info):
        # env_info is now a tuple of dicts
        env_info = [{k: v for k, v in dic.items() if k in self.aux_info} for dic in env_info]
        env_info = {k: [env_info[_][k] for _ in range(len(env_info))] for k in env_info[0].keys()}
        # env_info is now a dict of lists
        return env_info

    def fill_dictionaries(self, index, env_info, extra_predictions):
        for info in self.aux_info:
            dtype = torch.long if required_heads[info].startswith('multiclass') else torch.float
            self.collected_info[info][index] = torch.tensor(env_info[info], dtype=dtype, device=self.device)
            self.extra_predictions[info][index] = extra_predictions[info]

    def end_collection(self, exps):
        collected_info = dict()
        extra_predictions = dict()
        for info in self.aux_info:
            # T x P -> P x T -> P * T
            collected_info[info] = self.collected_info[info].transpose(0, 1).reshape(-1)
            if required_heads[info] == 'binary' or required_heads[info].startswith('continuous'):
                # T x P x 1 -> P x T x 1 -> P * T
                extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1)
            elif type(required_heads[info]) == int:
                # T x P x k -> P x T x k -> (P * T) x k
                k = required_heads[info]  # number of classes
                extra_predictions[info] = self.extra_predictions[info].transpose(0, 1).reshape(-1, k)
        # convert the dicts to DictLists, and add them to the exps DictList.
        exps.collected_info = DictList(collected_info)

        exps.extra_predictions = DictList(extra_predictions)