Dict observation space flattened

I want to use Dictionary observations in the model. However, RLlib seems to automatically flatten gym dict spaces. The custom_preprocessor settings has been deprecated. So I can’t find a way to prevent rllib from flattening the observation space.
Does anyone have an idea how I could prevent rllib from flattening the observation space?

You can still get to the original observation inside your custom model by accessing the input_dict’s “obs” key, no?


thank you for your answer.

No, the input_dict[‘obs’] contains only a flattened Tensor.
And also the make_model function receives only the flattened obs_space.

I am using the dreamer algorithm, if that helps.

Hmm, ok, could you try the following then inside your model code?

from ray.rllib.models.modelv2 import restore_original_dimensions

# inside your model's forward:
original_obs = restore_original_dimensions(input_dict["obs"], self.ob_space.original_space, "[tf|torch]")

The function seems to expect self.obs_space and not self.obs_space.original_space.
But with that change the code works. I can recover the original obs.

The problem is that the ‘forward’ function in the DreamerModel is called ‘policy’. That’s why the restore_original_dimensions is not automatically called

Got it. Makes sense. Thanks for reporting back @alex_koch !