smearle / control-pcgrl Goto Github PK

Expected Behavior

When running train_pcgrl multiagent.n_agents=2, the MultiAgentWrapper step method should execute without error.

Actual Behavior

When running train_pcgrl multiagent.n_agents=2, the environment is wrapped in a MultiAgentWrapper. When calling the step function of the wrapper, we call the step method of the parent class.

def step(self, action):
    obs, rew, done, info = {}, {}, {}, {}
    for k, v in action.items():
        self.unwrapped._rep.set_active_agent(k)
        obs_k, rew[k], done[k], info[k] = super().step(action={k: v}). # THIS LINE HERE
        obs.update(obs_k)
    done['__all__'] = np.all(list(done.values()))
    return obs, rew, done, info

The step method is expected to return the observations for a single agent, however, the parent's class observation space contains both agents. This raises the following exception:

AssertionError: The obs returned by the `step()` method observation keys is not same as the observation space keys, obs keys: ['agent_0'], space keys: ['agent_0', 'agent_1']

Environment

Ray 2.1.0
gym 0.26.2

Potential Fix

This implementation of the MultiAgentWrapper calls the step method of the parent class. Instead, we could call the step method of the unwrapped class. Note that this is not the base environment, this would just be one layer deeper:

    def step(self, action):
        obs, rew, done, info = {}, {}, {}, {}
        for k, v in action.items():
            self.unwrapped._rep.set_active_agent(k)
            obs_k, rew[k], done[k], info[k] = self.unwrapped.step(action={k: v}) ## CHANGED THIS LINE
            obs.update(obs_k)

        done['__all__'] = np.all(list(done.values()))

        return obs, rew, done, info

Note that since the representation is also wrapped, we can still pass in a action dictionary to the environment.