def __init__()

in mae_envs/wrappers/multi_agent.py [0:0]

• 19 lines of code
• 7 McCabe index (conditional complexity)

    def __init__(self, env, keys_self, keys_other, flatten=False):
        super().__init__(env)
        self.keys_self = sorted([k + '_self' for k in keys_self])
        self.keys_other = sorted(keys_other)
        self.flatten = flatten

        # Change observation space to look like a single agent observation space.
        # This makes constructing policies much easier
        if flatten:
            size_self = sum([np.prod(self.env.observation_space.spaces[k].shape[1:])
                             for k in self.keys_self + self.keys_other])
            self.observation_space = Dict(
                {'observation_self': Box(-np.inf, np.inf, (size_self,), np.float32)})
        else:
            size_self = sum([self.env.observation_space.spaces[k].shape[1]
                             for k in self.keys_self])
            obs_self = {'observation_self': Box(-np.inf, np.inf, (size_self,), np.float32)}
            obs_extern = {k: Box(-np.inf, np.inf, v.shape[1:], np.float32)
                          for k, v in self.observation_space.spaces.items()
                          if k in self.keys_other}
            obs_self.update(obs_extern)
            self.observation_space = Dict(obs_self)