[RLlib] Observation space with 2 dimensions not working with the new API stack #46631

I want to use an environment with an observation space of 2 dimensions with the new API stack but I’m unable to do so as the _get_encoder_config method of the Catalog class doesn’t have an encoder for such an environment. I believe that the issue is related to the error that is raised in the Catalog class:

ray\rllib\core\models\catalog.py line 368
elif (
                isinstance(observation_space, Box) and len(observation_space.shape) == 2
            ):
                # RLlib used to support 2D Box spaces by silently flattening them
                raise ValueError(
                    f"No default encoder config for obs space={observation_space},"
                    f" lstm={use_lstm} and attention={use_attention} found. 2D Box "
                    f"spaces are not supported. They should be either flattened to a "
                    f"1D Box space or enhanced to be a 3D box space."
                )

If this issue cannot be solved for now, what recommendations can I use to fix this issue locally? Note that the environment’s observation space cannot be changed. I also want to utilize the new API stack.
For more info about the issue, find it at [RLlib] Observation space with 2 dimensions not working with the new API stack · Issue #46631 · ray-project/ray · GitHub