[RLlib] Fix tree.flatten dict ordering bug: flatten_space([obs_space]) should produce same struct as tree.flatten([obs]).

[sven1977]

Fix tree.flatten dict ordering bug: flatten_space([obs_space]) should produce same struct as tree.flatten([obs]).

Simple reproduction of the issue:

from gym.spaces import Dict, Discrete
from ray.rllib.utils.space_utils import flatten_space
import tree  # dm_tree

space = Dict({"b": Discrete(200), "a": Discrete(1)})
flattened_spaces = flatten_space(space)
# flattened_spaces.spaces = [Discrete(200), Discrete(1)]  <- sorts dict keys by order of insertion!

flat_sample = tree.flatten(space.sample())
# flat_sample = [50, 0]  <- sorts dict keys by alphabet(!)

The above discrepancy could lead to mismatches (even unnoticed) when complex spaces are used, e.g. in our default ComplexInputNets (for tf and torch).

Why are these changes needed?

Related issue number

Checks

• I've run scripts/format.sh to lint the changes in this PR.
• I've included any doc changes needed for https://docs.ray.io/en/master/.
• I've made sure the tests are passing. Note that there might be a few flaky tests, see the recent failures at https://flakey-tests.ray.io/
• Testing Strategy
  • Unit tests
  • Release tests
  • This PR is not tested :(

[gjoliver]

awesome change.
thanks for the comments, and also fixing my dataset input bug.