Repository: google-deepmind/meltingpot
Issue: How to compute custom metrics during training ?
Root cause: Documentation Deficiency

Hello, 

In many of your papers you usually compute many metrics during the training process, i.e equality, sustainability, etc.
I am trying to compute these metrics for one of your substrates using RLlib. According to [This official tutorial](https://github.com/ray-project/ray/blob/3dda4e1d46a6cc33b11f3ed4cae3fbc88a29bda9/rllib/algorithms/callbacks.py), you can use custom callbacks to achieve that. 

Theoretically, the `on_episode_step` callback has a `base_env` parameter that stores the environment information, and you can simply access the environment data  using `obs, rewards, dones, infos, off_policy_actions = base_env.poll()`

In the case of _allelopathic_harvest_  some useful information is stored as WORLD observations, like in the case of WORLD.WHO_ZAPPED_WHO.  However, since the "WORLD" observations are deleted in [timestep_to_observations](https://github.com/deepmind/meltingpot/blob/28ba58b9bd0e138fd0a621e60f315c59b44d879c/examples/utils.py#L27) I can not use this data in the RLlib callback.

I tried to send data using the _info_  variable [in this script](https://github.com/deepmind/meltingpot/blob/28ba58b9bd0e138fd0a621e60f315c59b44d879c/examples/rllib/utils.py#L60) as show below. 

```
  def step(self, action):
    """See base class."""
    actions = [action[agent_id] for agent_id in self._ordered_agent_ids]
    timestep = self._env.step(actions)
    rewards = {
        agent_id: timestep.reward[index]
        for index, agent_id in enumerate(self._ordered_agent_ids)
    }
    done = {'__all__': True if timestep.last() else False}
    info = {"player_0": {"__common__ ": "test"}}

    observations = _timestep_to_observations(timestep)
    return observations, rewards, done, info

```
But in the callback side I just get the following output. `(RolloutWorker pid=21172) {0: {}}`. Additionally, I noticed that the rewards obtained in the callback side using  `base_env.poll()` are also empty.

Do you know how can we compute these metrics?
I am aware that you don't use RLlib internally in DeepMind, however I consider that this is the best place to ask. 

I think you can keep the world info in your buffer and calculate the metric. For agent execution, using world info is not allowed. 

> I think you can keep the world info in your buffer and calculate the metric. For agent execution, using world info is not allowed.

Hello and thank you!! 
What do you mean with in your buffer ? 

I managed to read the WORLD data on the RLlib callback by changing the [step method](https://github.com/deepmind/meltingpot/blob/5c8ac0f33f6286d792a2ce1b51482c414a15b67f/examples/rllib/utils.py#L51) to: 

```
  def step(self, action):
    """See base class."""
    actions = [action[agent_id] for agent_id in self._ordered_agent_ids]
    timestep = self._env.step(actions)

    rewards = {
        agent_id: timestep.reward[index]
        for index, agent_id in enumerate(self._ordered_agent_ids)
    }
    done = {'__all__': True if timestep.last() else False}
    info = _timestep_to_world_observations(timestep)

    observations = _timestep_to_observations(timestep)
    return observations, rewards, done, info
```
and added the following `_timestep_to_world_observations` function.
```

def _timestep_to_world_observations(timestep: dm_env.TimeStep):
  gym_observations = {}
  for index, observation in enumerate(timestep.observation):
    gym_observations[PLAYER_STR_FORMAT.format(index=index)] = {
        key: value for key, value in observation.items() if 'WORLD' in key
    }
  return gym_observations

```

In the Rllib side I gather the world data like this:

```
    def on_episode_step(self, worker: RolloutWorker, base_env: BaseEnv,
                        episode: MultiAgentEpisode, **kwargs):
        world_data = episode.last_info_for("player_0"))
```

However, _i)_ I am not sure if this is the best solution _ii)_ given that it needs to modify meltingpot's source code I don't know if it is worth to add it to the repository. 

> > I think you can keep the world info in your buffer and calculate the metric. For agent execution, using world info is not allowed.
> 
> Hello and thank you!! What do you mean with in your buffer ?

Hi, I mean in my code, I keep the WORLD info and collect them during training. I think DeepMind does not use RLLib in their project.  I think you can following the example here: https://github.com/ray-project/ray/blob/master/rllib/examples/custom_metrics_and_callbacks.py

I don't know how to integrate with RLLib, but there's this other thread where I explain how to get access to everything from the substrate without going via the RL loop: https://github.com/deepmind/meltingpot/issues/33

Hopefully that's useful.

Thank you !

I think that the code that I posted [here](https://github.com/deepmind/meltingpot/issues/45#issuecomment-1159237947) makes the job.

So I will close the Issue !

> Thank you !
> 
> I think that the code that I posted [here](https://github.com/deepmind/meltingpot/issues/45#issuecomment-1159237947) makes the job.
> 
> So I will close the Issue !

Great!!!