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?
Hi,
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]")
1 Like
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
1 Like
Got it. Makes sense. Thanks for reporting back @alex_koch !