Question

I am trying to build a dataset for the MiniHack environment. The environment has a fairly large dictionary space, with many keys, but only one level of nesting. I don't want to save all the keys, because it bloats the size of the dataset very much. I thought it could be done with StepDataCallback, like this:

class TTYStepDataCallback(minari.StepDataCallback):
    def __call__(self, env, obs, info, action=None, rew=None, terminated=None, truncated=None):
        tty_filter_keys = ["tty_chars", "tty_colors", "tty_cursor"]
        obs = {k: v for k, v in obs.items() if k in tty_filter_keys}

        step_data = super().__call__(env, obs, info, action, rew, terminated, truncated)
        
        return step_data

However, this will raise an error, as some keys from original observation space are missing. Why can't I just use the ones I want to save right away? I am logging data during distributed PPO training, which uses the other keys too, so they are needed for online training :(